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

楔子

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)])

最后,我们传递的列表里面的元素除了可以是 Series、列表,还可以是字典,举个例子。

import polars as pl

df = pl.DataFrame(
    [
        {"col1": 0, "col2": 3},
        {"col1": 2, "col2": 7},
        {"col1": 4, "col2": 5, "col3": 3},
    ]
)
print(df)
"""
shape: (3, 3)
┌──────┬──────┬──────┐
│ col1 ┆ col2 ┆ col3 │
│ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ i64  │
╞══════╪══════╪══════╡
│ 0    ┆ 3    ┆ null │
│ 2    ┆ 7    ┆ null │
│ 4    ┆ 5    ┆ 3    │
└──────┴──────┴──────┘
"""

列表里的字典中相同的字段会被归为一列(或者说一个字典就是一行数据),如果某个字段不存在,那么会被设置为空,在 Polars 里面空使用 null 来表示。

以上就是通过内置数据结构创建 DataFrame,还是比较简单的,和 Pandas 类似。大部分情况下,我们会传一个字典,或者包含多个字典的列表。

读取 CSV

读取 CSV 文件,Polars 有自己的快速实现,并且支持很多参数。当前有一个 girl.csv 文件,内容如下:

然后用 Polars 读取它。

import polars as pl

df = pl.read_csv("girl.csv")
print(df)
"""
shape: (3, 5)
┌─────┬─────────┬─────┬────────┬────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    │
╞═════╪═════════╪═════╪════════╪════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  │
└─────┴─────────┴─────┴────────┴────────┘
"""

值得一提的是,如果你在读取 CSV 的时候程序崩溃了,并给出如下警告。

那么你需要再安装一个包 pip install polars-lts-cpu,安装完之后 pl.read_csv 就没问题了。根据提示,说 CPU 处理器缺少两个功能:avx,fma,因此让我们安装 polars-lts-cpu,以获得更好的兼容性。目前这个问题存在于使用 M1 芯片、ARM 架构的 Mac 机器上,Linux 和 Windows 则没有问题。

好啦,目前 CSV 文件读取成功,下面看看 read_csv 的参数。


source

read_csv 的第一个参数,可以传一个表示文件路径的字符串,也可以传文件句柄、bytes 对象。

import polars as pl

# 文件路径
df1 = pl.read_csv("girl.csv")
# 文件句柄
with open("girl.csv", "rb") as f:
    df2 = pl.read_csv(f)
# 也可以传一个二进制字节流
with open("girl.csv", "rb") as f:
    df3 = pl.read_csv(f.read())

支持的格式很齐全,当你有一段 CSV 字节流时,可以直接传给 read_csv 函数。多说一句,在创建 df2 的时候会抛出警告,因为相比传一个文件句柄,Polars 更建议直接传一个文件路径。

以上就是 source 参数,而除了 source 之外还有很多其它参数,这些参数均要求以关键字参数的方式传递。


has_header

CSV 文件是否已经存在表头,默认为 True,会将第一行解析为表头,剩余行解析为数据。如果设置为 False,则表示 CSV 文件里面没有表头,全部是数据,此时 Polars 会自动生成表头。

import polars as pl

df1 = pl.read_csv("girl.csv")
print(df1)
"""
shape: (3, 5)
┌─────┬─────────┬─────┬────────┬────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    │
╞═════╪═════════╪═════╪════════╪════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  │
└─────┴─────────┴─────┴────────┴────────┘
"""
df2 = pl.read_csv("girl.csv", has_header=False)
print(df2)
"""
shape: (4, 5)
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│ column_1 ┆ column_2 ┆ column_3 ┆ column_4 ┆ column_5 │
│ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │
│ str      ┆ str      ┆ str      ┆ str      ┆ str      │
╞══════════╪══════════╪══════════╪══════════╪══════════╡
│ id       ┆ name     ┆ age      ┆ gender   ┆ length   │
│ 1        ┆ satori   ┆ 16       ┆ female   ┆ 155.3    │
│ 2        ┆ scarlet  ┆ 400      ┆ female   ┆ 145.9    │
│ 3        ┆ marisa   ┆ 18       ┆ female   ┆ 152.1    │
└──────────┴──────────┴──────────┴──────────┴──────────┘
"""

我们看到将 has_header 指定为 False 的时候,第一行也被解释为数据了,它原本是表头的。然后数据类型都变成了字符串,因为第一行数据都是字符串。


columns

如果 CSV 文件里面的列非常多,但我们只需要其中的几个,那么便可以通过 columns 参数选择指定的列。

import polars as pl

df1 = pl.read_csv("girl.csv", columns=["name", "length"])
print(df1)
"""
shape: (3, 2)
┌─────────┬────────┐
│ name    ┆ length │
│ ---     ┆ ---    │
│ str     ┆ f64    │
╞═════════╪════════╡
│ satori  ┆ 155.3  │
│ scarlet ┆ 145.9  │
│ marisa  ┆ 152.1  │
└─────────┴────────┘
"""
# 除了指定列名,也可以指定列的索引
df2 = pl.read_csv("girl.csv", columns=[1, 2])
print(df2)
"""
shape: (3, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ satori  ┆ 16  │
│ scarlet ┆ 400 │
│ marisa  ┆ 18  │
└─────────┴─────┘
"""

通过 columns 参数,便可避免读取无用数据。

另外关于 columns 参数,有一个需要注意的地方,看个例子。

import polars as pl

df = pl.read_csv("girl.csv", columns=["length", "name"])
print(df)
"""
shape: (3, 2)
┌─────────┬────────┐
│ name    ┆ length │
│ ---     ┆ ---    │
│ str     ┆ f64    │
╞═════════╪════════╡
│ satori  ┆ 155.3  │
│ scarlet ┆ 145.9  │
│ marisa  ┆ 152.1  │
└─────────┴────────┘
"""

columns 参数表示选择指定的列,但读取之后的字段顺序还是取决于 CSV。虽然这里 columns 指定的是 length、age,但读取之后 name 在 length 的前面,因为 CSV 里面字段 name 就在 length 的前面。


new_columns

如果你觉得 CSV 文件的列名不合适,想自己指定,那么便可以通过 new_columns 参数实现。

import polars as pl

df1 = pl.read_csv("girl.csv", new_columns=["编号", "姓名", "年龄", "性别", "身高"])
print(df1)
"""
┌──────┬─────────┬──────┬────────┬───────┐
│ 编号  ┆ 姓名    ┆ 年龄  ┆ 性别    ┆ 身高  │
│ ---  ┆ ---     ┆ ---  ┆ ---    ┆ ---   │
│ i64  ┆ str     ┆ i64  ┆ str    ┆ f64   │
╞══════╪═════════╪══════╪════════╪═══════╡
│ 1    ┆ satori  ┆ 16   ┆ female ┆ 155.3 │
│ 2    ┆ scarlet ┆ 400  ┆ female ┆ 145.9 │
│ 3    ┆ marisa  ┆ 18   ┆ female ┆ 152.1 │
└──────┴─────────┴──────┴────────┴───────┘
"""

# 如果 new_columns 的长度不够,那么剩余的列将保持原有的名称
df2 = pl.read_csv("girl.csv", new_columns=["编号", "姓名"])
print(df2)
"""
shape: (3, 5)
┌──────┬─────────┬─────┬────────┬────────┐
│ 编号  ┆ 姓名    ┆ age ┆ gender ┆ length │
│ ---  ┆ ---     ┆ --- ┆ ---    ┆ ---    │
│ i64  ┆ str     ┆ i64 ┆ str    ┆ f64    │
╞══════╪═════════╪═════╪════════╪════════╡
│ 1    ┆ satori  ┆ 16  ┆ female ┆ 155.3  │
│ 2    ┆ scarlet ┆ 400 ┆ female ┆ 145.9  │
│ 3    ┆ marisa  ┆ 18  ┆ female ┆ 152.1  │
└──────┴─────────┴─────┴────────┴────────┘
"""

columns 和 new_columns 可以结合使用,不指定 columns 会选择所有的列,指定 columns 可以选择出指定的列,然后通过 new_columns 重命名。

import polars as pl

df1 = pl.read_csv("girl.csv", columns=["id", "length"], new_columns=["ID", "LENGTH"])
print(df1)
"""
shape: (3, 2)
┌─────┬────────┐
│ ID  ┆ LENGTH │
│ --- ┆ ---    │
│ i64 ┆ f64    │
╞═════╪════════╡
│ 1   ┆ 155.3  │
│ 2   ┆ 145.9  │
│ 3   ┆ 152.1  │
└─────┴────────┘
"""

# 还是需要注意 columns 的顺序问题,虽然 columns 指定的是 ["length", "id"],但字段顺序是由 CSV 决定的
# 在 CSV 里面 id 在 length 的前面,所以选择之后,在 DataFrame 里面 id 依旧在 length 的前面
# 那么重命名为 ["LENGTH", "ID"] 就会出问题
df2 = pl.read_csv("girl.csv", columns=["length", "id"], new_columns=["LENGTH", "ID"])
print(df2)
"""
shape: (3, 2)
┌────────┬───────┐
│ LENGTH ┆ ID    │
│ ---    ┆ ---   │
│ i64    ┆ f64   │
╞════════╪═══════╡
│ 1      ┆ 155.3 │
│ 2      ┆ 145.9 │
│ 3      ┆ 152.1 │
└────────┴───────┘
"""
# 所以 new_columns 应该依旧是 ["ID", "LENGTH"]

以上就是 new_columns,不过它主要用在 CSV 没有表头的时候。因为 CSV 没有表头的话,我们会指定 has_header 为 False(否则会将第一行数据当做表头),让 Polars 自动生成表头,然后再通过 new_columns 重命名。


separator

分隔符,默认是逗号,这个比较简单,没什么可说的。


skip_rows

有时候数据文件不是从第一行开始的,因为一些用户可能会在开头写一些描述之类的,几行之后才是表头和数据。那么通过 skip_rows 参数可以跳过指定的行数,比如第三行是表头,就指定 skip_rows 为 2,跳过前两行。

import polars as pl

df = pl.read_csv("girl.csv", skip_rows=2)
print(df)
"""
shape: (1, 5)
┌─────┬─────────┬─────┬────────┬───────┐
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9 │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---   │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64   │
╞═════╪═════════╪═════╪════════╪═══════╡
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1 │
└─────┴─────────┴─────┴────────┴───────┘
"""

原有的表头和第一条数据被跳过去了,第二条数据(第三行)被当成了表头。


dtypes

Polars 在解析 CSV 时,会推断每一列的类型,而通过 dtypes 参数可以显式规定类型。意思是你别猜了,直接按我规定的类型解析就行。

import polars as pl

df = pl.read_csv("girl.csv", dtypes={"id": pl.Float64})
print(df)
"""
shape: (3, 5)
┌─────┬─────────┬─────┬────────┬────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    │
│ f64 ┆ str     ┆ i64 ┆ str    ┆ f64    │
╞═════╪═════════╪═════╪════════╪════════╡
│ 1.0 ┆ satori  ┆ 16  ┆ female ┆ 155.3  │
│ 2.0 ┆ scarlet ┆ 400 ┆ female ┆ 145.9  │
│ 3.0 ┆ marisa  ┆ 18  ┆ female ┆ 152.1  │
└─────┴─────────┴─────┴────────┴────────┘
"""

此时 id 就变成了浮点数,但如果指定的类型解析不了,那么 Polars 还是会继续推断,选择一个合适的类型。


null_values

将哪些数据判定为空数据,首先当分隔符的右边啥也没有,就是一个空数据,我们修改一下 CSV 文件:

id,name,age,gender,length
1,satori,16,female,155.3
2,scarlet,,female,145.9
3,marisa,18,female,152.1

我们看第二条数据,显然它的 age 字段就是空,因为逗号右边啥也没有。

import polars as pl

df = pl.read_csv("girl.csv")
print(df)
"""
shape: (3, 5)
┌─────┬─────────┬──────┬────────┬────────┐
│ id  ┆ name    ┆ age  ┆ gender ┆ length │
│ --- ┆ ---     ┆ ---  ┆ ---    ┆ ---    │
│ i64 ┆ str     ┆ i64  ┆ str    ┆ f64    │
╞═════╪═════════╪══════╪════════╪════════╡
│ 1   ┆ satori  ┆ 16   ┆ female ┆ 155.3  │
│ 2   ┆ scarlet ┆ null ┆ female ┆ 145.9  │
│ 3   ┆ marisa  ┆ 18   ┆ female ┆ 152.1  │
└─────┴─────────┴──────┴────────┴────────┘
"""

空数据读进来之后是个 null,这个 null 代表啥我们后续再说。但如果你希望某个具体的数据读进来之后也变成空,要怎么做呢?

import polars as pl

# 通过 null_values="16" 表示,所有值为 "16" 的 一律被解释为空字符串
df = pl.read_csv("girl.csv", null_values="16")
print(df)
"""
shape: (3, 5)
┌─────┬─────────┬──────┬────────┬────────┐
│ id  ┆ name    ┆ age  ┆ gender ┆ length │
│ --- ┆ ---     ┆ ---  ┆ ---    ┆ ---    │
│ i64 ┆ str     ┆ i64  ┆ str    ┆ f64    │
╞═════╪═════════╪══════╪════════╪════════╡
│ 1   ┆ satori  ┆ null ┆ female ┆ 155.3  │
│ 2   ┆ scarlet ┆ null ┆ female ┆ 145.9  │
│ 3   ┆ marisa  ┆ 18   ┆ female ┆ 152.1  │
└─────┴─────────┴──────┴────────┴────────┘
"""

这里可能有人好奇,age 字段不是整数吗?为啥是字符串 "16" 呢?首先 CSV 只是一个普通的纯文本,字段类型是 Polars 解析数据之后推断出来的,在解析之前数据都被视为字符串,而 null_values 就是在此时完成的替换。

此时不仅原有的空数据被替换成了 null,"16" 也被换成了 null。另外 null_values 还可以是一个列表,支持接收多个字符串。

import polars as pl

df = pl.read_csv("girl.csv", null_values=["16", "2", "145.9"])
print(df)
"""
shape: (3, 5)
┌──────┬─────────┬──────┬────────┬────────┐
│ id   ┆ name    ┆ age  ┆ gender ┆ length │
│ ---  ┆ ---     ┆ ---  ┆ ---    ┆ ---    │
│ i64  ┆ str     ┆ i64  ┆ str    ┆ f64    │
╞══════╪═════════╪══════╪════════╪════════╡
│ 1    ┆ satori  ┆ null ┆ female ┆ 155.3  │
│ null ┆ scarlet ┆ null ┆ female ┆ null   │
│ 3    ┆ marisa  ┆ 18   ┆ female ┆ 152.1  │
└──────┴─────────┴──────┴────────┴────────┘
"""

所有在 null_values 里面的值都会被替换成空,不管哪一列。那么问题来了,如果你希望只替换指定的列,其它列不受影响,该怎么做呢?

import polars as pl

df = pl.read_csv("girl.csv", null_values={"id": "16"})
print(df)
"""
shape: (3, 5)
┌─────┬─────────┬──────┬────────┬────────┐
│ id  ┆ name    ┆ age  ┆ gender ┆ length │
│ --- ┆ ---     ┆ ---  ┆ ---    ┆ ---    │
│ i64 ┆ str     ┆ i64  ┆ str    ┆ f64    │
╞═════╪═════════╪══════╪════════╪════════╡
│ 1   ┆ satori  ┆ 16   ┆ female ┆ 155.3  │
│ 2   ┆ scarlet ┆ null ┆ female ┆ 145.9  │
│ 3   ┆ marisa  ┆ 18   ┆ female ┆ 152.1  │
└─────┴─────────┴──────┴────────┴────────┘
"""

只对列 id 里面的 16 进行替换,其它列则不影响,当然这里没有 id 为 16 的数据。因此传递字典,可以让不同的列里面的不同的值被替换。另外需要注意:当指定为字典时,每一列只能指定一个被替换的值,比如下面是不合法的。

null_values={"id": ["16", "2"]}

所以 null_values 只能接收以下三种类型的值:

  • str
  • List[str]
  • Dict[str, str]

try_parse_dates

是否解析日期,默认为 False,表示不解析。如果指定为 True,那么符合日期格式的字符串会被推断出来,从而解析成日期类型。若解析失败,依旧保持 pl.String 类型。

我们修改一下 girl.csv。

id,name,age,gender,length,join_time
1,satori,16,female,155.3,1998-12-11 12:43:18
2,scarlet,400,female,145.9,1997-08-21 00:00:00
3,marisa,18,female,152.1,2005-06-18 07:22:37

然后读取文件。

import polars as pl

df1 = pl.read_csv("girl.csv")
print(df1)
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    ┆ str                 │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

df2 = pl.read_csv("girl.csv", try_parse_dates=True)
print(df2)
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    ┆ datetime[μs]        │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

我们看到 join_time 被解析成日期格式了,当然啦,即使不用 try_parse_dates 参数也可以解析成日期。还记得前面的 dtype 参数吗?

import polars as pl

df = pl.read_csv("girl.csv", dtypes={"join_time": pl.Datetime})
print(df)
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    ┆ datetime[μs]        │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

像整数和浮点数可以自动推断出来,但日期如果不显示指定的话,会被解析成字符串。这时候可以通过 dtype 参数告诉 Polars,将 join_time 字段按照日期来解析。所以个人觉得 dtype 参数比 try_parse_dates 更合适一些,因为比起让 Polars 去猜,我们告诉 Polars 会更好一些,当然具体选择哪一种取决于你自己。


n_threads

读取 CSV 时使用多少个线程,默认使用本地 CPU 的逻辑核心数。这个参数很强大,可以并行读取,而 Pandas 只能使用单个核。


infer_schema_length

我们说过,Polars 在将数据读取进来之后会推断每一列的类型,而 infer_schema_length 则表示用于推断数据类型所需读取的最大行数,默认是 100。如果推断失败,比如 100 行之后推断出某个字段是 pl.Int64,但后续又发现该字段还包含了 pl.Float64 类型的值,那么会增加行数重新推断。

如果设置为 0,那么表示不推断,所有列都被解析为 pl.String。如果设置为 None,则将所有数据全部读取进来之后,再推断类型,此时是最准确的,但速度也会稍慢(相对来说)。

import polars as pl

df = pl.read_csv("girl.csv", infer_schema_length=0)
print(df)
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ str ┆ str     ┆ str ┆ str    ┆ str    ┆ str                 │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

将 infer_schema_length 设置为 0,所有列都被解析成了字符串。大部分情况下,该参数我们都不需要指定。


batch_size

一次性往缓冲区读入多少行,默认是 8192。


n_rows

如果 CSV 有很多行,但我们只需要指定数量的行,那么可以通过 n_rows 指定要读取的行数。n_rows 默认为 None,表示全部读取,如果你想只读取前 1w 行,那么就将 n_rows 指定为 10000 即可。

但要注意的是,在多线程情况下,不能严格遵守上限 n_rows,也就是读取的行数可能会超过 n_rows(但不会太多)。


encoding

处理文本时使用的编码,默认是 utf8,还可以指定 utf8-lossy,表示有损模式,有损模式意味着无效的 utf8 值会被替换为 �。


low_memory

是否通过牺牲性能来减少内存压力,默认为 False。

以上就是 read_csv 函数的一些参数,当然还有几个不常用的我们没有说,有兴趣可以自己了解一下。

然后读取 CSV 文件除了 read_csv 函数之外,还有 scan_csv,这两个函数的用法和参数都是一样的。但区别是 read_csv 会立即加载整个 CSV 文件到内存中,并返回一个完成的 DataFrame,当数据集可以被完整地载入内存时 read_csv 非常适合。

scan_csv 函数采用一种懒加载的方式来处理数据,它并不立即加载整个文件,而是创建一个可用于查询的虚拟 DataFrame。这种方法允许 Polars 的查询优化器进行更有效的数据处理,例如下推过滤条件和列投影到文件扫描级别。这意味着如果你只对数据的一小部分感兴趣,Polars 可能只会加载和处理那一小部分,而不是整个文件。

因此 scan_csv 更适合处理大型数据集,尤其是在内存容量有限的情况下,它可以显著减少内存使用和提高性能。

读取 Excel

说完了读取 CSV,再来看看读取 Excel,Excel 也是一种非常常见的外部文件。但是从性能的角度来看,在存储数据时不建议使用 Excel,而是使用 Parquet 或 CSV,原因有以下几点。

1)性能

Excel 文件是一种复杂的文件格式,包含了单元格格式、公式、元数据等信息,这使得读写操作相对较慢。而 Parquet 是一种列式存储格式,非常适合于大型数据集的高效读写操作。CSV 是一种简单的行文本格式,虽然没有 Parquet 那么高效,但通常也比 Excel 快。

2)内存使用

读写 Excel 文件通常比 Parquet 或 CSV 文件占用更多内存,这对于处理大数据集来说可能是一个问题。

3)数据规模和复杂性

对于小型数据集和那些需要利用 Excel 的特定功能(如格式化单元格、使用公式等)的应用场景,Excel 仍然是一个不错的选择。但对于大型数据集和主要用于数据分析的场景,Parquet 和 CSV 提供了更好的性能和可扩展性。

因此对于需要高性能处理的大型数据集,推荐使用 Parquet 或 CSV 格式。不过就目前来说,Excel 使用的依旧很广泛,所以 Polars 提供了 read_excel 函数。但需要注意,Polars 原生是不支持读取 Excel 的,它底层会使用其它的外部库进行读取,然后生成 DataFrame。

# 可选的外部库有三种,分别是 xlsx2csv openpyxl pyxlsb
pip install xlsx2csv openpyxl pyxlsb

然后我们来读取 Excel。

import polars as pl

# Excel 内部可以有多个 sheet,通过 sheet_id 指定要读取第几个 sheet
# 注意:sheet_id 从 1 开始
df = pl.read_excel("girl.xlsx", sheet_id=1)
print(df)
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    ┆ str                 │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

这里我将刚才 girl.csv 里面的内容复制了一份,放到了 girl.xlsx 里面,我们看到读取是没问题的。

再来说一下 sheet_id 参数,Excel 内部可以有多个 sheet,通过 sheet_id 可以指定要读取第几个 sheet。如果不指定 sheet_id,那么默认读取第一个。另外 read_excel 还支持同时读取多个 sheet,只需要给 sheet_id 传一个列表即可。

import polars as pl

# 我们这里只有一个 sheet
# 此时会返回一个字典,key 是 sheet 的名称,value 是对应的 DataFrame
df_dict = pl.read_excel("girl.xlsx", sheet_id=[1])
print(df_dict.__class__)  # <class 'dict'>
# 每个 sheet 都有一个名称,默认是 "Sheet1", "Sheet2", "Sheet3", ...
print(df_dict["Sheet1"])
"""
shape: (3, 6)
┌─────┬─────────┬─────┬────────┬────────┬─────────────────────┐
│ id  ┆ name    ┆ age ┆ gender ┆ length ┆ join_time           │
│ --- ┆ ---     ┆ --- ┆ ---    ┆ ---    ┆ ---                 │
│ i64 ┆ str     ┆ i64 ┆ str    ┆ f64    ┆ str                 │
╞═════╪═════════╪═════╪════════╪════════╪═════════════════════╡
│ 1   ┆ satori  ┆ 16  ┆ female ┆ 155.3  ┆ 1998-12-11 12:43:18 │
│ 2   ┆ scarlet ┆ 400 ┆ female ┆ 145.9  ┆ 1997-08-21 00:00:00 │
│ 3   ┆ marisa  ┆ 18  ┆ female ┆ 152.1  ┆ 2005-06-18 07:22:37 │
└─────┴─────────┴─────┴────────┴────────┴─────────────────────┘
"""

问题来了,如果我们想读取全部的 sheet 该怎么办呢?很简单,直接给 sheet_id 传个 0 即可,此时就会读取全部的 sheet。

然后选择 sheet 除了通过 sheet_id 参数之外,还可以通过 sheet_name。通过指定 sheet 的名字来筛选会更方便一些,特别是在 sheet 非常多的时候。

import polars as pl

# 需要注意的是,sheet_id 和 sheet_name 不可以同时指定
df = pl.read_excel("girl.xlsx", sheet_name="Sheet1")
print(df)  # <class 'dict'>
# 和 sheet_id 一样,sheet_name 也可以接收一个列表,同时读取多个 sheet
df_dict = pl.read_excel("girl.xlsx", sheet_name=["Sheet1"])
print(df_dict["Sheet1"])

补充:在读取 Excel 的时候,第一个参数可以是一个路径,也可以是一个文件句柄或者 bytes 对象。

然后 read_excel 还有一个 engine 参数,用于指定读取 Excel 所使用的外部库。

import polars as pl

# engine 有三种选择,分别是 xlsx2csv、openpyxl、pyxlsb,默认是 xlsx2csv
df = pl.read_excel("girl.xlsx", engine="xlsx2csv")
print(df)

这里推荐使用 xlsx2csv。

读取 Parquet

Parquet 是一种列式存储格式,它具备非常多的优点。首先与行式存储(如 CSV)相比,在列式存储中,同一列的数据被连续存储在一起,这种存储方式对分析型查询非常有效,因为它允许快速读取、处理和聚合特定的列数据,而不需要加载整个数据集。

另外列式存储对数据压缩更为友好,事实上列式存储和数据压缩是伴生的。由于同一列中的数据类型相同,且往往具有一定的相关性,因此可以应用更高效的压缩算法。这不仅减少了存储空间的需求,还加快了数据的读取速度。

然后是 Polars DataFrame 和 Parquet 的布局相似性,Polars DataFrame 的内存布局在许多方面反映了 Parquet 文件在磁盘上的布局,这意味着将 Polars DataFrame 转换为 Parquet 文件(反之亦然)时,所需的数据转换非常少,从而加快了读写速度。

由于这些特性,Parquet 格式特别适合大型数据集的高效处理,尤其是在需要对特定列进行频繁读取和分析的场景中。此外 Parquet 还是大数据生态系统中常用的文件格式,非常适用于需要高性能数据读写、有效的存储利用和快速查询访问的应用场景。因此当处理大规模数据集且对性能有较高要求时,使用 Parquet 格式而不是 CSV 或其它行存储格式,通常是更佳的选择。

说完了 Parquet 文件的优点,来看看 Polars 如何读取 Parquet 文件。

df = pl.read_parquet("docs/data/path.parquet")

只需要调用 read_parquet 函数即可,非常简单。关于 read_parquet 还有很多其它参数,这里就不赘述了,如果你了解 Parquet 文件格式,那么这些参数直接就见名知意了。如果不熟悉的话,可以上官网查看一下,而且事实上我们用第一个参数就足够了。

然后除了 read_parquet 之外,还有 scan_parquet,作用和 scan_csv 类似。

读取数据库

读取数据库应该是最常见的了,大部分情况下我们都需要从数据库读取数据,Polars 也提供了相应的函数支持。

from sqlalchemy import create_engine
import polars as pl

query = "SELECT * FROM user_to_role"
engine = create_engine("mysql+pymysql://")
# 通过 read_database 函数即可读取数据库
# 第一个参数是 SQL 语句,第二个参数是引擎或者链接
df = pl.read_database(query, engine)
print(df)
"""
shape: (9, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 1       ┆ 1       │
│ 2       ┆ 1       │
│ 3       ┆ 1       │
│ 4       ┆ 1       │
│ 7       ┆ 1       │
│ 8       ┆ 1       │
│ 9       ┆ 1       │
│ 10      ┆ 1       │
│ 6       ┆ 2       │
└─────────┴─────────┘
"""

非常方便,类似于 Pandas 的 read_sql,该函数还接收其它参数,来介绍一下。


iter_batches 和 batch_size

如果数据集过大,那么可以将 iter_batches 指定为 True,此时会返回一个生成器,每迭代一次返回一批数据。通过这种方式可以将大型结果集分批加载到内存中,而无需一次性加载整个结果集。至于每一批返回多少条数据,则由 batch_size 指定,该参数接收一个整数,比如 8192,那么每迭代一次,就会返回 8192 条数据组成的 DataFrame。

from sqlalchemy import create_engine
import polars as pl

query = "SELECT * FROM user_to_role"
engine = create_engine("mysql+pymysql://")
df_iter = pl.read_database(query, engine, iter_batches=True, batch_size=4)
print(df_iter)
"""
<generator object ConnectionExecutor._from_rows.<locals>.<genexpr> at 0x7f8b08d7ad60>
"""

print(df_iter.__next__())
"""
shape: (4, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 1       ┆ 1       │
│ 2       ┆ 1       │
│ 3       ┆ 1       │
│ 4       ┆ 1       │
└─────────┴─────────┘
"""

print(df_iter.__next__())
"""
shape: (4, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 7       ┆ 1       │
│ 8       ┆ 1       │
│ 9       ┆ 1       │
│ 10      ┆ 1       │
└─────────┴─────────┘
"""

print(df_iter.__next__())
"""
shape: (1, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 6       ┆ 2       │
└─────────┴─────────┘
"""

try:
    print(df_iter.__next__())
except StopIteration:
    print("迭代完毕")
"""
迭代完毕
"""

当数据表非常大时,便可以使用这种方式。


schema_overrides

数据库的表的每一列都有严格的类型,但如果你觉得这些类型不准确的话,那么也可以通过 schema_overrides 参数将其覆盖掉。

from sqlalchemy import create_engine
import polars as pl

query = "SELECT * FROM user_to_role"
engine = create_engine("mysql+pymysql://")
# 这里我们将 user_id 改成了字符串,当然我们改成字符串反而是不对的,这里只是演示这个功能
df = pl.read_database(query, engine, schema_overrides={"user_id": pl.String})
print(df)
"""
shape: (9, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ str     ┆ i64     │
╞═════════╪═════════╡
│ 1       ┆ 1       │
│ 2       ┆ 1       │
│ 3       ┆ 1       │
│ 4       ┆ 1       │
│ 7       ┆ 1       │
│ 8       ┆ 1       │
│ 9       ┆ 1       │
│ 10      ┆ 1       │
│ 6       ┆ 2       │
└─────────┴─────────┘
"""

对于数据库来说,这个参数不太常用。


execute_options

可以用来传递一些执行时的参数,举个例子。

from sqlalchemy import create_engine
import polars as pl

query = "SELECT * FROM user_to_role WHERE user_id > :user_id"
engine = create_engine("mysql+pymysql://")
# query 里面有一个占位符,它的值可以通过 execute_options 指定
# Polars 会通过 execute_options["parameters"]["user_id"] 拿到指定的值,并将占位符替换掉
df = pl.read_database(query, engine,
                      execute_options={"parameters": {"user_id": 7}})
print(df)
"""
shape: (3, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 8       ┆ 1       │
│ 9       ┆ 1       │
│ 10      ┆ 1       │
└─────────┴─────────┘
"""

这个参数也不常用,query 里面占位符的值一般都会提前指定好,比如:

from sqlalchemy import create_engine, text
import polars as pl

query = text("SELECT * FROM user_to_role WHERE user_id > :user_id").bindparams(user_id=7)
engine = create_engine("mysql+pymysql://")
df = pl.read_database(query, engine)
print(df)
"""
shape: (3, 2)
┌─────────┬─────────┐
│ user_id ┆ role_id │
│ ---     ┆ ---     │
│ i64     ┆ i64     │
╞═════════╪═════════╡
│ 8       ┆ 1       │
│ 9       ┆ 1       │
│ 10      ┆ 1       │
└─────────┴─────────┘
"""

当然具体使用哪一种取决于你自己的喜好,总之 Polars 可以从各种数据库里面读取想要的数据。

读取云存储

内置的数据结构、文件,以及数据库,这些数据源已经满足日常工作中的大部分场景。但 Polars 还支持从 AWS S3、Azure Blob、Google Cloud 等云存储中读取数据,并且 API 是相同的。当然,要从这些云存储中读取数据,需要安装的额外的依赖包。

pip install fsspec s3fs adlfs gcsfs

然后 Polars 可以直接从云存储中读取 CSV、IPC、Parquet 等文件,就仿佛这些文件是在本地一样。

import polars as pl

source = "s3://bucket/*.parquet"

df = pl.read_parquet(source)

那么问题来了,位于 S3 等云存储上的文件一般不能直接访问,而是需要传递 Access key 等进行验证,那么这些用于验证的参数该怎么传呢?很简单,read_csv 和 read_parquet 内部都有一个 storage_options 参数,它们是专门用来认证的。

import polars as pl

source = "s3://bucket/*.parquet"

storage_options = {
    "aws_access_key_id": "<secret>",
    "aws_secret_access_key": "<secret>",
    "aws_region": "us-east-1",
}
df = pl.read_parquet(source, storage_options=storage_options)  

还是挺方便的,并且 CSV、Parquet 等文件可以使用通配符批量读取,无论是本地文件,还是云存储上的文件。

但我们知道,对于像 S3 这种云存储来说,它是按是扫描的数据量来收费的。read_parquet 会将所有文件都下载到内存中(通过 fsspec 管理内部下载),并创建一个 DataFrame(相当于把每个文件读取出来后,合并在一块了),这是比较耗时的(而且还耗钱)。所以我们还可以使用 scan_parquet 函数,它的用法、参数和 read_parquet 是一样的,区别就在于 read_parquet 属于即席查询,返回的是 DataFrame,而 scan_parquet 属于延迟查询,返回的是 LazyFrame。

通过创建带有谓词和投影下推的延迟查询,查询优化器显著减少需要下载的数据量,从而减少成本。CSV 文件也是同理,如果文件非常多,但是只有一部分满足查询条件,那么便可以使用 scan_csv。

读取 BigQuery

BigQuery 是 Google 推出的一项服务,该服务可以让开发者使用 Google 的架构来运行 SQL 语句,对超级大的数据库表进行操作。BigQuery 允许用户上传他们的超大量数据并进行交互式分析,从而不必投资建立自己的数据中心。据 Google 表示,BigQuery 引擎可以快速扫描高达 70TB 未经压缩处理的数据,并且马上得到分析结果。

但 BigQuery 存在的目的并不是为了替换掉数据库,它不适合实时的高频交易,BigQuery 的宗旨是对大量的数据集(比如数十亿行)进行分析和计算,从而得到相应的结果。

Polars 也支持从 BigQuery 查询数据,但需要安装一个依赖项。

pip install google-cloud-bigquery

然后便可以让 BigQuery 执行查询,获取数据了。

import polars as pl
from google.cloud import bigquery

client = bigquery.Client()

# Perform a query.
QUERY = (
    'SELECT name FROM `bigquery-public-data.usa_names.usa_1910_2013` '
    'WHERE state = "TX" '
    'LIMIT 100')
query_job = client.query(QUERY)  # API request
rows = query_job.result()  # Waits for query to finish

df = pl.from_arrow(rows.to_arrow())

总的来说,BigQuery 用的不多,至少我从工作到现在没有用过。

导出 DataFrame 到指定数据源

说完数据的读取,再来看看 DataFrame 的导出。说实话,做数据处理总共就三步:

  • 读取数据
  • 对数据做处理
  • 导出处理后的数据

我们将数据读入进来之后,得到一个 DataFrame 或 LazyFrame,然后做各种各样的处理,具体视需求而定。当数据处理完之后,还是一个 DataFrame 的形式,所以最终还要将 DataFrame 导出到指定的数据源。那么下面就来看看 DataFrame 的导出函数。

导出为内置数据结构

首先是将 DataFrame 导出为列表、字典等数据结构。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
# as_series 默认为 True,如果为 True,那么导出的字典的 value 就是 Series
# 这里我们一般指定为 False
print(df.to_dict(as_series=False))
"""
{'name': ['satori', 'scarlet', 'marisa'], 
 'age': [16, 400, 18], 
 'length': [155.3, 145.9, 152.1]}
"""
# 如果我们返回一个列表,每行数据都是列表里的一个字典,要怎么做呢?
print(df.to_dicts())
"""
[{'name': 'satori', 'age': 16, 'length': 155.3}, 
 {'name': 'scarlet', 'age': 400, 'length': 145.9},
 {'name': 'marisa', 'age': 18, 'length': 152.1}]
"""

比较简单,关于内置数据结构,导出为这两种就足够了。

导出为 CSV

然后是导出为 CSV。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
# 参数可以是一个文件路径,也可以是一个以写模式打开的文件句柄
df.write_csv("girl.csv")

然后该方法还接收很多其它参数,这些参数都要以关键字参数的方式指定,下面分别介绍一下。

  • include_bom:输出的 CSV 是否包含 UTF-8 BOM,默认为 False。

  • include_header:输出的 CSV 是否包含头部,默认为 True。

  • separator:分隔符,默认是逗号。

  • line_terminator:行终止符,也就是一行的结尾,默认是 \n。

  • quote_char:如果单元格的值本身就包含分隔符,那么要如何区分呢?这个时候就会使用 quote_char 将该单元格包起来,默认是 "。

  • batch_size:每个线程处理的行数,默认是 1024,如果行数少于 1024,那么就只有一个线程。

  • datetime_format:接收一个自定义的 format 字符串,比如 datetime_format="%Y年%m月%d日 %H时%M分%S秒",那么 2024-01-12 18:25:22 写入到 CSV 中就是 "2024年01月12日 18时25分22秒"。当然,你也可以将日期转成指定的字符串之后,再写入到 CSV。

  • date_format:和 datetime_format 类似,但针对的是 pl.Date 类型。

  • time_format:针对 pl.Time 类型。

  • float_precision:写入到 CSV 中的浮点数要保留多少位,比如保留 5 位,那么多了截断,少了补零。

  • null_value:null 值以什么形式写到 CSV,默认是空字符串。你也可以指定为其它的,比如 "NA",这样 null 值写入到 CSV 文件时就是 NA。

  • quote_style:将单元格的值使用引号包起来所对应的策略,可选的值为 "necessary"、"always"、"non_numeric"、"never"。

    • necessary:默认选项,只有在必要时才会使用引号包起来,比如单元格本身包含 "、分隔符、行终止符等。
    • always:总是使用引号将单元格包起来。
    • never:任何单元格都不使用引号包起来,即使这可能会导致 CSV 无效。
    • non_numeric:将所有不是数值的单元格使用引号包起来。

我们举个例子:

from datetime import datetime
import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 2, 3],
        "b": [1.414, 3.1415926, 2.71],
        "c": [None, None, "军统六哥"],
        "d": [datetime(2023, 11, 24, 12, 22, 56),
              datetime(2023, 8, 16, 22, 9, 33),
              datetime(2023, 2, 17, 9, 23, 40)]
    }
)
df.write_csv(
    "file.csv",
    # 浮点数保留 5 位小数
    float_precision=5,
    # null 值以字符串 "我是空" 的形式写入
    null_value="我是空",
    # 日期也转换一下
    datetime_format="%Y 年 %m 月 %d 日 %H 时 %M 分 %S 秒"
)

看一下生成的数据。

以上就是 CSV 的导出,还是比较简单的。

导出为 Excel

再来看看导出为 Excel,导出为 Excel 需要安装一个外部依赖 xlsxwriter,不安装会报错。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
# 写入 Excel,第一个参数可以是文件路径
# 也可以是一个以二进制写模式打开的文件句柄,包含 io.BytesIO()
# 第二个参数表示 sheet 的名称,默认为 "Sheet1"
df.write_excel("girl.xlsx", worksheet="新的 shit")

导出成功,看一下生成的 Excel。

结果没有任何问题,然后 write_excel 在生成 Excel 的时候还支持很多其它参数,这些参数都要以关键字参数的方式指定,下面分别介绍一下。


position

把数据看做一个矩形,那么可以通过 position 指定矩形的起始位置,比如我们指定为 "B4",那么生成的 Excel 就是下面这个样子。

另外 position 还可以接收一个元组,比如 (3, 1),表示左上角位于第 4 行、第 2 列的位置,和字符串 "B4" 是等价的。不过我们一般还是习惯指定成字符串,当然啦,position 这个参数本身就不怎么常用。


table_style

生成的 Excel 表的风格,举个例子。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
df.write_excel("girl.xlsx", table_style="Table Style Medium 4")

此时生成的 Excel 如下:

具有可以设置哪些 style,可以网上搜索,不过作为开发一般只负责导出数据,一般很少会做样式这些花里胡哨的。

  • float_precision:和 write_csv 中的 float_precision 用法一致,默认为 3。
  • include_header:导出的数据是否包含表头,默认为 True。
  • autofit:生成的 Excel 是否针对内容的长度,对单元格进行拉伸,默认为 False。此时单元格内容比较长的话,会被隐藏掉,需要手动拉伸单元格。

以上就是 write_excel 的一些常用的参数,当然还有很多其它参数,主要都是跟样式相关的,这里不赘述了,有兴趣可以自己查看,注释非常详细。

导出为 Parquet

再来看看如何导出为 Parquet 文件。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
# 写入 Excel,第一个参数可以是文件路径
# 也可以是一个以二进制写模式打开的文件句柄,包含 io.BytesIO()
df.write_parquet("girl.parquet")

执行没有问题,成功写入 Parquet 文件,下面来看看其它参数。

  • compression:数据压缩方式,支持 lz4、uncompressed、snappy、gzip、lzo、brotli、zstd,默认是 "zstd"。

  • compression_level:压缩级别,值越高,压缩率越高,生成的文件就越小,当然耗时也会更长。

    • "gzip":0 ~ 10;
    • "brotli":0 ~ 11;
    • "zstd":1 ~ 22;
  • statistics:是否将统计信息写入 Parquet 文件的表头中,这需要额外的计算。

  • row_group_size:行分组的大小,默认是 512 ^ 2 行。

  • data_page_size:数据页的大小,以字节为单位,默认是 1024 ^ 2 字节。

大部分情况下,我们只需要指定第一个参数即可。

导出为数据库表数据

除了导出到文件,还可以导出到指定的数据库表中。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1]}
)
# 写入数据库表,第一个参数是表名,第二个参数是创建引擎的 URI
# if_table_exists 参数有三种选择:fail、replace、append
#     如果指定为 fail,当数据表存在时报错
#     如果指定为 replace,当数据表存在时会先删除再重建,然后导入数据
#     如果指定为 append,当数据表存在时会追加
# 最后是 engine 参数,可以选择 sqlalchemy 或 adbc,表示用哪个库创建引擎,默认是 sqlalchemy
df.write_database("girl", "mysql+pymysql://",
                  if_table_exists="append",
                  engine="sqlalchemy")

关于导出到数据库表中,我个人更习惯先导出为包含多个字典的列表,然后手动创建引擎,最后将数据导进去。

如果引擎使用 sqlalchemy 创建,那么要求必须安装 pandas,所以不建议使用该函数。

关于数据的导出暂时就说这么多,在工作中应该是够用了。

基于 DataFrame 处理数据

说完了读取数据、DataFrame 的导出之后,我们来看最重要的一个步骤,也就是数据处理,或者说数据清洗等等。在这个步骤中,我们可以改变每一列的值,基于某一列进行筛选、聚合,以及做一些更复杂的操作等等。Polars 提供了非常丰富的 API,可以对 DataFrame 做各种变换,而且这些 API 的速度也是极快的。

好了, 下面我们来介绍这些 API。

select 方法,选择指定字段

一个 DataFrame 可以有很多列,通过 select 方法可以选择指定的列。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)
# 选择 name、length 两列
print(df.select("name", "length"))
"""
shape: (3, 2)
┌─────────┬────────┐
│ name    ┆ length │
│ ---     ┆ ---    │
│ str     ┆ f64    │
╞═════════╪════════╡
│ satori  ┆ 155.3  │
│ scarlet ┆ 145.9  │
│ marisa  ┆ 152.1  │
└─────────┴────────┘
"""
# 也可以传一个列表,这两者是等价的
print(df.select(["name", "length"]))
"""
shape: (3, 2)
┌─────────┬────────┐
│ name    ┆ length │
│ ---     ┆ ---    │
│ str     ┆ f64    │
╞═════════╪════════╡
│ satori  ┆ 155.3  │
│ scarlet ┆ 145.9  │
│ marisa  ┆ 152.1  │
└─────────┴────────┘
"""

select 在选择的时候,还可以改变某一列的值。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)
print(df.select("name", pl.col("age") + 1))
"""
shape: (3, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ satori  ┆ 17  │
│ scarlet ┆ 401 │
│ marisa  ┆ 19  │
└─────────┴─────┘
"""

获取列的时候,可以传 "列名" 或者 pl.col("列名"),但如果要改变某一列则只能使用 pl.col("列名")。像上面的代码,如果写成 "age" + 1 显然就报错了。

然后 pl.col() 里面也可以传多个列名:

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)
# 在筛选的时候,以下几种方式都是等价的,都是筛选 age 和 length 两列
"""
df.select("age", "length")
df.select(["age", "length"])
df.select(pl.col("age"), pl.col("length"))
df.select([pl.col("age"), pl.col("length")])
df.select(pl.col("age", "length"))
df.select([pl.col("age", "length")])
"""
# 不过为了清晰,个人建议不用在外面再嵌套一层列表了
# 然后筛选时不直接指定字段,而是统一使用 pl.col(字段) 的形式


# 表示将 length 和 age 字段都加 1
# 等价于 df.select(pl.col("length") + 1, pl.col("age") + 1)
print(df.select(pl.col("length", "age") + 1))
"""
shape: (3, 2)
┌────────┬─────┐
│ length ┆ age │
│ ---    ┆ --- │
│ f64    ┆ i64 │
╞════════╪═════╡
│ 156.3  ┆ 17  │
│ 146.9  ┆ 401 │
│ 153.1  ┆ 19  │
└────────┴─────┘
"""

需要注意的是,df.select() 会返回一个新的 DataFrame,不会影响原有的 df。

然后在筛选的时候,还可以指定通配符 *,表示筛选所有的列,因此非常适合拷贝一份 DataFrame。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)
print(df.select(pl.col("*")))
"""
shape: (3, 4)
┌─────────┬─────┬────────┬────────┐
│ name    ┆ age ┆ length ┆ salary │
│ ---     ┆ --- ┆ ---    ┆ ---    │
│ str     ┆ i64 ┆ f64    ┆ i64    │
╞═════════╪═════╪════════╪════════╡
│ satori  ┆ 16  ┆ 155.3  ┆ 6000   │
│ scarlet ┆ 400 ┆ 145.9  ┆ 7500   │
│ marisa  ┆ 18  ┆ 152.1  ┆ 5000   │
└─────────┴─────┴────────┴────────┘
"""

在筛选的时候,还可以按照字段类型进行筛选,关于字段类型都有哪些,一会再总结。

import polars as pl
import polars.selectors as cs

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)

# 选择 pl.Int64 和 pl.Float64 类型的字段
print(df.select(pl.col(pl.Int64, pl.Float64)))
"""
shape: (3, 3)
┌─────┬────────┬────────┐
│ age ┆ length ┆ salary │
│ --- ┆ ---    ┆ ---    │
│ i64 ┆ f64    ┆ i64    │
╞═════╪════════╪════════╡
│ 16  ┆ 155.3  ┆ 6000   │
│ 400 ┆ 145.9  ┆ 7500   │
│ 18  ┆ 152.1  ┆ 5000   │
└─────┴────────┴────────┘
"""
# 选择 Int64 类型的字段,并将字段的值加 1
# 选择 Float64 类型的字段,并将字段的值加 1.5
print(df.select(pl.col(pl.Int64) + 1, pl.col(pl.Float64) + 1.5))
"""
shape: (3, 3)
┌─────┬────────┬────────┐
│ age ┆ salary ┆ length │
│ --- ┆ ---    ┆ ---    │
│ i64 ┆ i64    ┆ f64    │
╞═════╪════════╪════════╡
│ 17  ┆ 6001   ┆ 156.8  │
│ 401 ┆ 7501   ┆ 147.4  │
│ 19  ┆ 5001   ┆ 153.6  │
└─────┴────────┴────────┘
"""

Polars 对类型要求比较严格,pl.Int32 和 pl.Int64 是不同的类型。

按照字段名筛选的时候,可以直接往 df.select 方法中传入字段名,或者传入 pl.col(字段名)。但按照类型筛选时,则必须传 pl.col(类型)。

就目前来讲,关于选择指定字段已经很差不多了,但 Polars 还专门提供了选择器,用于实现更复杂的筛选。

import polars as pl
import polars.selectors as cs

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)
# 选择全部字段
print(df.select(cs.all()))

# 选择第一个字段和最后一个字段
print(df.select(cs.first(), cs.last()))

# 选择名称以 "na" 开头的字段
print(df.select(cs.starts_with("na")))
# 选择名称不以 "na" 开头的字段,这些选择器均支持取反操作
print(df.select(~cs.starts_with("na")))

# 选择名称以 "e" 结尾的字段
print(df.select(cs.ends_with("e")))
# 选择名称不以 "e" 结尾的字段
print(df.select(~cs.ends_with("e")))

# 选择类型包含 "g" 的字段
print(df.select(cs.contains("g")))
# 选择类型不包含 "g" 的字段
print(df.select(~cs.contains("g")))

# 选择类型为字符串的字段
# 我们说这些选择器均支持取反,所以 ~cs.string() 就表示选择类型不是字符串的字段
print(df.select(cs.string()))

# 选择类型为整数的字段,不限制精度
print(df.select(cs.integer()))

# 选择类型为浮点数的字段
print(df.select(cs.float()))

# 选择类型为 decimal 的字段
print(df.select(cs.decimal()))

# 选择类型为布尔值的字段
print(df.select(cs.boolean()))

# 选择字段类型为 bytes 的字段
print(df.select(cs.binary()))

# 选择类型为 datetime 的字段
print(df.select(cs.datetime()))
# 选择类型为 date 的字段
print(df.select(cs.date()))
# 选择类型为 time 的字段
print(df.select(cs.time()))
# 选择类型为 timedelta 的字段
print(df.select(cs.duration()))

# 选择类型为 categorical 的字段
print(df.select(cs.categorical()))

# 按照字段名选择字段,等价于 df.select(pl.col("name"), pl.col("age"))
print(df.select(cs.by_name("name", "age")))
# 选择 "name" 字段之外的其它字段,当要筛选的字段非常多时,这种做法就很方便
print(df.select(~cs.by_name("name")))

# 按照类型选择字段,等价于 df.select(pl.col(pl.Int64), pl.col(pl.Float64))
print(df.select(cs.by_dtype(pl.Int64, pl.Float64)))

关于字段选择就说到这里,选择方式可以说非常的丰富,具体使用哪一种看个人喜好以及场景。不过大部分情况下,我们都会直接基于字段做筛选。

filter 方法,对数据进行过滤

select 方法是选择指定的字段,也就是选择指定的列,而 filter 则是选择指定的行,比如选择 age 字段大于 18 的行。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)

# 选择 age >= 18 的行
print(df.filter(pl.col("age") > 18))
"""
shape: (1, 4)
┌─────────┬─────┬────────┬────────┐
│ name    ┆ age ┆ length ┆ salary │
│ ---     ┆ --- ┆ ---    ┆ ---    │
│ str     ┆ i64 ┆ f64    ┆ i64    │
╞═════════╪═════╪════════╪════════╡
│ scarlet ┆ 400 ┆ 145.9  ┆ 7500   │
└─────────┴─────┴────────┴────────┘
"""

select 和 filter 方法也可以结合起来使用,比如先选择指定的字段,再进行过滤;或者先对数据进行过滤,再选择指定的字段。

import polars as pl

df = pl.DataFrame(
    {"name": ["satori", "scarlet", "marisa"],
     "age": [16, 400, 18],
     "length": [155.3, 145.9, 152.1],
     "salary": [6000, 7500, 5000]}
)

# 选择 age > 18 的行,然后选择 "name" 和 "age" 列
print(df.filter(pl.col("age") > 18).select(pl.col("name", "age")))
"""
shape: (1, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ scarlet ┆ 400 │
└─────────┴─────┘
"""

# 先选择 "name" 和 "age" 列,然后再选择 age > 18 的行
# 注意:在选择的时候我们给 "age" 字段加了 1
print(df.select(pl.col("name"), pl.col("age") + 1).filter(pl.col("age") > 18))
"""
shape: (2, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ scarlet ┆ 401 │
│ marisa  ┆ 19  │
└─────────┴─────┘
"""
# 无论是 select 还是 filter,都会返回一个 DataFrame
# 所以上面的 filter 是对 select 之后的 DataFrame 进行过滤
# 如果再反过来,此时是先对数据行进行过滤,再将 "age" 字段加 1,所以结果是不一样的
print(df.filter(pl.col("age") > 18).select(pl.col("name"), pl.col("age") + 1))
"""
shape: (1, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ scarlet ┆ 401 │
└─────────┴─────┘
"""

然后关于过滤,需要补充一下 Pandas 和 Polars 之间的区别。首先无论是 Pandas 还是 Polars,它们的 DataFrame 的每一列都是一个 Series,而多个 Series 并排站在一起就组成了 DataFrame。

s = pl.Series()
# 在 Pandas 里面,我们会使用 df["age"] > 18 这种形式,会先手动拿到相应的 Series,然后比较
# 而在 Polars 里面,则使用 pl.col("age") > 18 这种形式,可以认为这个过程只是定义了一个操作
# 当具体执行时,内部会自动获取对应的 Series,然后进行比较。这样做就带来了一个好处,比如:
# df.filter(pl.col("age") > 18).filter(pl.col("age") > 20),显然这个操作是冗余的
# 那么 Polars 就会将这两步优化为一步,只按照 pl.col("age") > 20 进行过滤
print(df.filter(pl.col("age") > 18))

所以 filter 可以链式调用,因为它返回的是 DataFrame,调用多少次都可以。当然啦,相比多次调用,每次调用指定一个过滤条件,我们更多的是将过滤条件指定完整,然后只调用一次(当然具体视情况而定)。那么问题来了,我们都可以按照哪些条件进行过滤呢?毫无疑问,只要是 Series 对象支持的方法,都可以通过 pl.col("字段") 进行调用、对数据做过滤。

我们随便举几个例子:

import polars as pl

df = pl.DataFrame(
    {"num": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
     "alpha": ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]}
)

# 选择 num 字段小于 5 的
print(df.filter(pl.col("num") < 5))
"""
shape: (4, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ a     │
│ 2   ┆ b     │
│ 3   ┆ c     │
│ 4   ┆ d     │
└─────┴───────┘
"""
# 选则 num 字段为偶数的
print(df.filter(pl.col("num") & 1 == 0))
"""
shape: (5, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 2   ┆ b     │
│ 4   ┆ d     │
│ 6   ┆ f     │
│ 8   ┆ h     │
│ 10  ┆ j     │
└─────┴───────┘
"""
# 选择 num 字段乘以 2 之后,大于 16 的
print(df.filter(pl.col("num") * 2 > 16))
"""
shape: (2, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 9   ┆ i     │
│ 10  ┆ j     │
└─────┴───────┘
"""

当字段为数值时,可以直接比较,这一点和 Pandas 是一样的。如果字段的类型是字符串,我要如何按照是否包含某个子串、是否小于指定长度、是否大写等条件进行过滤呢?

通过 str 属性会返回一个 ExprStringNameSpace 对象,它里面有非常多的方法,足以满足工作中的所有需求。这些方法都见名知意,源码内部也有非常详细的注释,这里就不赘述了。

from polars.expr.string import ExprStringNameSpace

直接进入到源码中查看即可。

如果是日期类型,那么就获取 .dt 属性,会返回一个 polars.expr.datetime.ExprDateTimeNameSpace 对象。方法同样见名知意,并且注释详细。

Polars 对列表也有很丰富的支持,如果某一列的数据类型是 pl.List,那么便可以获取 list 属性,进行操作。

如果有 Pandas 的经验,这些方法应该都很简单,不仅注释详细,而且类型注解也很完善。有兴趣的话,可以看看这些方法,不用刻意去记,在工作中多用几次就熟悉了。当然,我们后续会实际操作,来对这些内容进行完善。

一些基础概念

这一部分来说一说 Polars 里面的基础概念, 理解了它们,才能更好地使用 Polars。

类型

首先是类型,DataFrame 的每一列都是一个 Series,每个 Series 包含名称、数据、类型,那么类型都有哪些呢?

import polars as pl

# 整数
pl.UInt8
pl.UInt16
pl.UInt32
pl.UInt64
pl.Int8
pl.Int16
pl.Int32
pl.Int64
# 浮点数
pl.Float32
pl.Float64
# 字符串
pl.String
# Categorical
pl.Categorical
# 布尔值
pl.Boolean
# bytes
pl.Binary
# 列表
pl.List
# datetime
pl.Datetime
# date
pl.Date
# time
pl.Time
# timedelta
pl.Duration
# object
pl.Object
# Enum
pl.Enum

我们举例说明:

from datetime import datetime
import polars as pl

df = pl.DataFrame(
    {
        "integer": [1, 2],
        "floating": [3.14, 2.71],
        "string": ["hello", "world"],
        "boolean": [True, False],
        "binary": [b"ping", b"pong"],
        "list": [[1, 2], [3, 4]],
        "datetime": [datetime(2033, 12, 15, 22, 33, 44)] * 2
    }
)
print(df)
"""
shape: (2, 7)
┌─────────┬──────────┬────────┬─────────┬───────────────┬───────────┬─────────────────────┐
│ integer ┆ floating ┆ string ┆ boolean ┆ binary        ┆ list      ┆ datetime            │
│ ---     ┆ ---      ┆ ---    ┆ ---     ┆ ---           ┆ ---       ┆ ---                 │
│ i64     ┆ f64      ┆ str    ┆ bool    ┆ binary        ┆ list[i64] ┆ datetime[μs]        │
╞═════════╪══════════╪════════╪═════════╪═══════════════╪═══════════╪═════════════════════╡
│ 1       ┆ 3.14     ┆ hello  ┆ true    ┆ [binary data] ┆ [1, 2]    ┆ 2033-12-15 22:33:44 │
│ 2       ┆ 2.71     ┆ world  ┆ false   ┆ [binary data] ┆ [3, 4]    ┆ 2033-12-15 22:33:44 │
└─────────┴──────────┴────────┴─────────┴───────────────┴───────────┴─────────────────────┘
"""

注意:在使用类型时,我们要用 Polars 的类型。

然后需要说一说浮点数,Polars 提供的 Float32 和 Float64 遵循 IEEE754 规则,但也有一些例外:

  • 任何的 NaN 之间都是相等的,并且大于任何非 NaN 值。

  • 在进行数学运算或数据处理操作(如排序或分组)时,对零的符号(正零或负零)和 NaN(非数字)的处理不保证任何特定行为,比如:

    • 在某些数学运算中,零可能具有正负号(+0 或 -0),但在这些操作中,不保证保留零的原始符号。例如,在排序或分组操作中,所有的零可能被规范化为正零(+0)。
    • 对于 NaN(非数字),操作不保证保持其原始的有效载荷(即 NaN 背后的特定编码信息)。所有的 NaN 在操作过程中被统一处理为标准形式的正 NaN,而不包含任何特定的有效载荷信息。

这种规范化处理主要是为了进行高效的相等性检查,例如在排序或分组数据时,统一零的符号和 NaN 的形式可以简化以及加快比较操作。

数据结构

Polars 提供了两种关键的数据结构,分别是 Series 和 DataFrame。Series 是一个一维数据结构,由名称、数据、类型三部分组成,一个 Series 里面的元素的类型是相同的。

import polars as pl

# 如果不指定 dtype,那么默认为 Int64
s = pl.Series("age", [18, 16, 22], dtype=pl.UInt8)
print(s)
"""
shape: (3,)
Series: 'age' [i64]
[
    18
    16
    22
]
"""

DataFrame 是一个由多个 Series 组成的二维数据结构,可以在 DataFrame 上执行的操作与在 SQL 的查询中执行的操作非常相似,可以 GROUP BY、JOIN、PIVOT 等,还可以定义自定义函数。

import polars as pl

df = pl.DataFrame(
    {
        "num": [1, 2, 3, 4, 5],
        "alpha": ["a", "b", "c", "d", "e"]
    }
)
print(df)
"""
shape: (5, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ a     │
│ 2   ┆ b     │
│ 3   ┆ c     │
│ 4   ┆ d     │
│ 5   ┆ e     │
└─────┴───────┘
"""
# 选择前 n 行
print(df.head(3))
"""
shape: (3, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ a     │
│ 2   ┆ b     │
│ 3   ┆ c     │
└─────┴───────┘
"""
# 选择后 n 行
print(df.tail(2))
"""
shape: (2, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 4   ┆ d     │
│ 5   ┆ e     │
└─────┴───────┘
"""
# 采样,随机抽取两个,这个方法还是很有用的
# 当你需要对大批量数据做类似均值统计时,不妨先抽样,如果计算结果和预期大差不差,再统计全部数据
print(df.sample(2))
"""
shape: (2, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 2   ┆ b     │
│ 5   ┆ e     │
└─────┴───────┘
"""
# 还可以按照比例抽样,抽 60% 的数据
# with_replacement:抽样时数据是否有放回,默认 False。如果指定为 True,那么数据可能会有重复
# shuffle:抽样前数据是否打乱,默认 False
# seed:随机种子
print(df.sample(fraction=0.6, with_replacement=False, shuffle=False, seed=666))
"""
shape: (3, 2)
┌─────┬───────┐
│ num ┆ alpha │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 4   ┆ d     │
│ 2   ┆ b     │
│ 1   ┆ a     │
└─────┴───────┘
"""
# 如果你想查看某个 DataFrame 的统计信息,还可以调用 describe 方法
print(df.describe())
"""
┌────────────┬──────────┬───────┐
│ describe   ┆ num      ┆ alpha │
│ ---        ┆ ---      ┆ ---   │
│ str        ┆ f64      ┆ str   │
╞════════════╪══════════╪═══════╡
│ count      ┆ 5.0      ┆ 5     │
│ null_count ┆ 0.0      ┆ 0     │
│ mean       ┆ 3.0      ┆ null  │
│ std        ┆ 1.581139 ┆ null  │
│ min        ┆ 1.0      ┆ a     │
│ 25%        ┆ 2.0      ┆ null  │
│ 50%        ┆ 3.0      ┆ null  │
│ 75%        ┆ 4.0      ┆ null  │
│ max        ┆ 5.0      ┆ e     │
└────────────┴──────────┴───────┘
"""

这些概念属实简单了,我们继续往下看。

一些基础操作

像 df.select()、df.filter() 等等在 Polars 里面被称之为表达式,而在表达式中可以使用各种数学运算符。

import numpy as np
import polars as pl

df = pl.DataFrame(
    {
        "nrs": [1, 2, 3, None, 5],
        "names": ["foo", "ham", "spam", "egg", None],
        "random": np.random.rand(5),
        "groups": ["A", "A", "B", "C", "B"],
    }
)
print(df)
"""
shape: (5, 4)
┌──────┬───────┬──────────┬────────┐
│ nrs  ┆ names ┆ random   ┆ groups │
│ ---  ┆ ---   ┆ ---      ┆ ---    │
│ i64  ┆ str   ┆ f64      ┆ str    │
╞══════╪═══════╪══════════╪════════╡
│ 1    ┆ foo   ┆ 0.197785 ┆ A      │
│ 2    ┆ ham   ┆ 0.1587   ┆ A      │
│ 3    ┆ spam  ┆ 0.777543 ┆ B      │
│ null ┆ egg   ┆ 0.90939  ┆ C      │
│ 5    ┆ null  ┆ 0.644459 ┆ B      │
└──────┴───────┴──────────┴────────┘
"""
df_numerical = df.select(
    # 执行运算之后,字段名保持不变,但可以调用 alias 起一个别名
    (pl.col("nrs") + 5).alias("nrs + 5"),
    (pl.col("nrs") - 5).alias("nrs - 5"),
    (pl.col("nrs") * pl.col("random")).alias("nrs * random"),
    (pl.col("nrs") / pl.col("random")).alias("nrs / random"),
)
print(df_numerical)
"""
shape: (5, 4)
┌─────────┬─────────┬──────────────┬──────────────┐
│ nrs + 5 ┆ nrs - 5 ┆ nrs * random ┆ nrs / random │
│ ---     ┆ ---     ┆ ---          ┆ ---          │
│ i64     ┆ i64     ┆ f64          ┆ f64          │
╞═════════╪═════════╪══════════════╪══════════════╡
│ 6       ┆ -4      ┆ 0.197785     ┆ 5.055988     │
│ 7       ┆ -3      ┆ 0.3174       ┆ 12.602395    │
│ 8       ┆ -2      ┆ 2.332628     ┆ 3.858309     │
│ null    ┆ null    ┆ null         ┆ null         │
│ 10      ┆ 0       ┆ 3.222297     ┆ 7.758442     │
└─────────┴─────────┴──────────────┴──────────────┘
"""
df_logical = df.select(
    (pl.col("nrs") > 1).alias("nrs > 1"),
    (pl.col("random") <= 0.5).alias("random <= .5"),
    (pl.col("nrs") != 1).alias("nrs != 1"),
    (pl.col("nrs") == 1).alias("nrs == 1"),
    ((pl.col("random") <= 0.5) & (pl.col("nrs") > 1)).alias("and_expr"),  # and
    ((pl.col("random") <= 0.5) | (pl.col("nrs") > 1)).alias("or_expr"),  # or
)
print(df_logical)
"""
shape: (5, 6)
┌─────────┬──────────────┬──────────┬──────────┬──────────┬─────────┐
│ nrs > 1 ┆ random <= .5 ┆ nrs != 1 ┆ nrs == 1 ┆ and_expr ┆ or_expr │
│ ---     ┆ ---          ┆ ---      ┆ ---      ┆ ---      ┆ ---     │
│ bool    ┆ bool         ┆ bool     ┆ bool     ┆ bool     ┆ bool    │
╞═════════╪══════════════╪══════════╪══════════╪══════════╪═════════╡
│ false   ┆ true         ┆ false    ┆ true     ┆ false    ┆ true    │
│ true    ┆ true         ┆ true     ┆ false    ┆ true     ┆ true    │
│ true    ┆ false        ┆ true     ┆ false    ┆ false    ┆ true    │
│ null    ┆ false        ┆ null     ┆ null     ┆ false    ┆ null    │
│ true    ┆ false        ┆ true     ┆ false    ┆ false    ┆ true    │
└─────────┴──────────────┴──────────┴──────────┴──────────┴─────────┘
"""

我们注意到,在筛选列的时候,可以对列执行相关操作改变列数据。比如将列 nrs 加上 5,再比如执行比较操作,此时得到的就是一组布尔值。如果这些比较操作放在了 filter 方法里,那么就是对数据进行过滤。

整个过程应该很好理解,如果是 Pandas,那么代码逻辑会类似于 df["nrs"] = df["nrs"] + 5 这种形式,这会改变原有的 DataFrame。而 Polars 的 df.select 方法是将感兴趣的列筛选出来,相当于拷贝了一份,在拷贝之后的数据上做操作。

需要注意的是,在筛选时列不可以重复。

import numpy as np
import polars as pl

df = pl.DataFrame(
    {
        "nrs": [1, 2, 3, None, 5],
        "names": ["foo", "ham", "spam", "egg", None],
        "random": np.random.rand(5),
        "groups": ["A", "A", "B", "C", "B"],
    }
)

try:
    df_numerical = df.select(
        df.select(pl.col("nrs") + 5, pl.col("nrs") - 5)
    )
except Exception as e:
    print(e)
"""
column with name 'nrs' has more than one occurrences
"""    

# 表达式在运算的时候,列名不会改变,因此 "nrs" 发生了重复
# 这个时候可以再调用 .alias() 起一个别名

我们也可以统计不重复的值的数量。

import numpy as np
import polars as pl

df = pl.DataFrame(
    {
        "groups": ["A", "A", "B", "C", "B"],
    }
)

df = df.select(
    # n_unique 是精确统计,approx_n_unique 会采用 HyperLogLog++ 算法做近似统计,但效果足够好
    pl.col("groups").n_unique().alias("unique"),
    pl.approx_n_unique("groups").alias("unique_approx"),
)
print(df)
"""
shape: (1, 2)
┌────────┬───────────────┐
│ unique ┆ unique_approx │
│ ---    ┆ ---           │
│ u32    ┆ u32           │
╞════════╪═══════════════╡
│ 3      ┆ 3             │
└────────┴───────────────┘
"""

Polars 还在表达式中支持类似 if else,只不过使用的是 when、then 和 else 语法。谓词放在 when 子句中,当它的值为 True 时应用 then 表达式,否则应用 else 表达式。

import polars as pl

df = pl.DataFrame(
    {
        "nrs": [1, 2, 3, None, 5],
    }
)
df = df.select(pl.col("nrs"), pl.when(pl.col("nrs") > 2).then(True).otherwise(False).alias("conditional"))
print(df)
"""
shape: (5, 2)
┌──────┬─────────────┐
│ nrs  ┆ conditional │
│ ---  ┆ ---         │
│ i64  ┆ bool        │
╞══════╪═════════════╡
│ 1    ┆ false       │
│ 2    ┆ false       │
│ 3    ┆ true        │
│ null ┆ false       │
│ 5    ┆ true        │
└──────┴─────────────┘
"""

应该很好理解,我们看 Polars 有点像函数式编程,因为 Rust 这门语言就是这么设计的。

强制类型转换

使用 cast() 方法可以进行强制类型转换,将列的数据类型转换为新类型,Polars 使用 Arrow 来管理内存中的数据,并依赖于 Rust 实现的计算内核来进行转换。

cast 方法除了接收一个新类型之外,还包含一个 strict 参数,用于决定当 Polars 遇到无法从源数据类型转换为目标数据类型时的行为。默认情况下,strict=True,这意味着 Polar 将抛出一个错误,通知用户转换失败,并提供有关无法强制转换的值的详细信息。另一方面,如果 strict=False,任何不能转换为目标类型的值都将悄悄地转换为空。

import polars as pl
import polars.selectors as cs

df = pl.DataFrame(
    {
        "integers": [1, 2, 3, 4, 5],
        "floats": [4.0, 5.0, 6.0, 7.0, 8.0],
        "floats_with_decimal": [4.532, 5.5, 6.5, 7.5, 8.5],
    }
)
print(df)
"""
shape: (5, 3)
┌──────────┬────────┬─────────────────────┐
│ integers ┆ floats ┆ floats_with_decimal │
│ ---      ┆ ---    ┆ ---                 │
│ i64      ┆ f64    ┆ f64                 │
╞══════════╪════════╪═════════════════════╡
│ 1        ┆ 4.0    ┆ 4.532               │
│ 2        ┆ 5.0    ┆ 5.5                 │
│ 3        ┆ 6.0    ┆ 6.5                 │
│ 4        ┆ 7.0    ┆ 7.5                 │
│ 5        ┆ 8.0    ┆ 8.5                 │
└──────────┴────────┴─────────────────────┘
"""
out = df.select(
    pl.col("integers").cast(pl.Float32).alias("integers_as_floats"),
    pl.col("floats").cast(pl.Int32).alias("floats_as_integers"),
    pl.col("floats_with_decimal").cast(pl.Int32).alias("floats_with_decimal_as_integers"),
)
print(out)
"""
shape: (5, 3)
┌────────────────────┬────────────────────┬─────────────────────────────────┐
│ integers_as_floats ┆ floats_as_integers ┆ floats_with_decimal_as_integers │
│ ---                ┆ ---                ┆ ---                             │
│ f32                ┆ i32                ┆ i32                             │
╞════════════════════╪════════════════════╪═════════════════════════════════╡
│ 1.0                ┆ 4                  ┆ 4                               │
│ 2.0                ┆ 5                  ┆ 5                               │
│ 3.0                ┆ 6                  ┆ 6                               │
│ 4.0                ┆ 7                  ┆ 7                               │
│ 5.0                ┆ 8                  ┆ 8                               │
└────────────────────┴────────────────────┴─────────────────────────────────┘
"""

减少内存占用也可以通过修改分配给元素的位数来实现,比如从 Int64 转换到 Int16、从 Float64 转换到Float32。

import polars as pl
import polars.selectors as cs

df = pl.DataFrame(
    {
        "integers": [1, 2, 3, 4, 5],
        "floats": [4.0, 5.0, 6.0, 7.0, 8.0],
        "floats_with_decimal": [4.532, 5.5, 6.5, 7.5, 8.5],
    }
)

out = df.select(
    pl.col("integers").cast(pl.Int16).alias("integers_smallfootprint"),
    pl.col("floats").cast(pl.Float32).alias("floats_smallfootprint"),
)
print(out)
"""
shape: (5, 2)
┌─────────────────────────┬───────────────────────┐
│ integers_smallfootprint ┆ floats_smallfootprint │
│ ---                     ┆ ---                   │
│ i16                     ┆ f32                   │
╞═════════════════════════╪═══════════════════════╡
│ 1                       ┆ 4.0                   │
│ 2                       ┆ 5.0                   │
│ 3                       ┆ 6.0                   │
│ 4                       ┆ 7.0                   │
│ 5                       ┆ 8.0                   │
└─────────────────────────┴───────────────────────┘
"""

在执行向下转换时,确保所选择的位数足以容纳列中的最大值和最小值。例如 32 位有符号整数(Int32)可以处理 -2147483648 到 +2147483647 范围内的整数,而使用 Int8 可以处理 -128 到 127 之间的整数。尝试转换为太小的 DataType 将导致 Polars 抛出 ComputeError,因为该操作不支持。但如果将 strict 参数设置为 True,那么当溢出(无法转化)时会转成 null。

以上就是类型转换,你可以试试其它类型。

聚合操作

然后再来看看聚合操作,在 Polars 中你可以聚合任意的字段,没有任何限制。

import polars as pl

df = pl.DataFrame(
    [
        {"language": "Python", "framework": "FastAPI"},
        {"language": "Python", "framework": "Sanic"},
        {"language": "Python", "framework": "Blacksheep"},
        {"language": "Python", "framework": "Flask"},
        {"language": "Golang", "framework": "Gin"},
        {"language": "Golang", "framework": "Beego"},
        {"language": "Golang", "framework": "Iris"},
        {"language": "Rust", "framework": "Axum"},
        {"language": "Rust", "framework": "Tokio"},
    ]
)
# 按照 language 进行分组
out = df.group_by("language").agg(
    # 获取每组的元素个数
    pl.col("framework").count().alias("count"),
    # 获取每组的第一个元素
    pl.col("framework").first().alias("first"),
    # 获取每组的最后一个元素
    pl.col("framework").last().alias("last"),
    # 获取每组的最大值(字符串会比较字典序),这里求最大值没什么意义
    # pl.col("framework").max(),
    # 当然还可以求和,基本上你能想到的任何操作都可以实现
    # pl.col("framework").sum(),
)
print(out)
"""
shape: (3, 4)
┌──────────┬───────┬─────────┬───────┐
│ language ┆ count ┆ first   ┆ last  │
│ ---      ┆ ---   ┆ ---     ┆ ---   │
│ str      ┆ u32   ┆ str     ┆ str   │
╞══════════╪═══════╪═════════╪═══════╡
│ Python   ┆ 4     ┆ FastAPI ┆ Flask │
│ Golang   ┆ 3     ┆ Gin     ┆ Iris  │
│ Rust     ┆ 2     ┆ Axum    ┆ Tokio │
└──────────┴───────┴─────────┴───────┘
"""

如果我想在 group_by 的基础上再进行过滤要怎么办呢?很简单,调用一次 filter 即可。

out = df.group_by("language").agg(
    pl.col("framework").count().alias("count"),
    pl.col("framework").first().alias("first"),
    pl.col("framework").last().alias("last"),
).filter(pl.col("count") > 2)
print(out)
"""
shape: (2, 4)
┌──────────┬───────┬─────────┬───────┐
│ language ┆ count ┆ first   ┆ last  │
│ ---      ┆ ---   ┆ ---     ┆ ---   │
│ str      ┆ u32   ┆ str     ┆ str   │
╞══════════╪═══════╪═════════╪═══════╡
│ Golang   ┆ 3     ┆ Gin     ┆ Iris  │
│ Python   ┆ 4     ┆ FastAPI ┆ Flask │
└──────────┴───────┴─────────┴───────┘
"""

这些操作都是链式的,每一步都返回一个 DataFrame,我们是在聚合的结果上调用 filter,那么自然就是对聚合后的结果进行过滤。聚合后的 DataFrame 是有 count 字段的,因此我们可以基于 count 字段进行过滤。

未完待续

本文参考自:

  • https://docs.pola.rs/user-guide/
  • https://blog.jetbrains.com/zh-hans/dataspell/2023/09/polars-vs-pandas-what-s-the-difference/
posted @ 2024-01-12 00:04  古明地盆  阅读(1345)  评论(0编辑  收藏  举报