Pandas-为什么 Polars 比 Pandas 使用更少的内存

目录

为什么 Polars 比 Pandas 使用更少的内存

使用 Pandas 处理大量数据可能会很困难；很容易耗尽内存，导致程序变慢甚至崩溃。Polars 数据框库是一个潜在的解决方案。

虽然 Polars 主要以比 Pandas 运行更快而闻名，但如果使用得当，它有时也可以显著减少内存使用。特别是，某些在 Pandas 中需要手动完成的技术可以在 Polars 中自动完成，从而让你在处理大型数据集时使用更少的内存——并且减少你的工作量！

当然，这需要你使用正确的 Polars API。而且它并不能解决所有问题，即使它确实让你的生活更轻松。

在本文中，我们将：

1. 看看如何通过一些工作来优化 Pandas 的内存使用。
2. 看看 Polars 如何在某些情况下自动应用这些技术。
3. 指出至少一些你需要手动干预以减少内存使用的方式。

---

一个例子：从简单的 Pandas 到内存优化的 Pandas

为了帮助理解 Polars 如何比 Pandas 减少内存使用，我们将从一个具体的例子开始，并在 Pandas 中实现它。我们将查看波士顿地区交通管理局（MBTA）记录的公交路线时间，并尝试找出特别慢的公交路线。

我们将使用 2022 年的数据，可在此处获取，数据以 CSV 文件形式提供。每个月的数据是一个约 300MB 的 CSV 文件；我们将查看 5 月的数据。

首先，我们将使用一个简单的 Pandas 实现来实现一个特定的查询，然后使用一个更优化但仍然基于 Pandas 的实现。在下一节中，我们将切换到 Polars。

第一步：更高效的内存表示和文件格式

以下是数据的一个示例，省略了一些列：

service_date	route_id	direction_id	standard_type	scheduled	scheduled_headway	headway
2022-05-01	“01”	“Inbound”	“Schedule”	“1900-01-01 06:05:00.000”	NA	NA
2022-05-01	“01”	“Inbound”	“Schedule”	“1900-01-01 06:25:00.000”	NA	NA
2022-05-01	“01”	“Inbound”	“Headway”	“1900-01-01 06:25:00.000”	“1200”	“841”
2022-05-01	“01”	“Inbound”	“Schedule”	“1900-01-01 06:29:00.000”	NA	NA
2022-05-01	“01”	“Inbound”	“Schedule”	“1900-01-01 06:30:00.000”	NA	NA

Inbound 和 Outbound 表示前往或离开波士顿的方向，波士顿也被称为 “The Hub”。

首先要注意的是，许多列可以使用更高效的数据类型来表示，而不会丢失任何信息。

• 将 service_date、actual 和 scheduled_headway 从字符串转换为时间戳。
• 将 route_id、direction_id 和 standard_type 从字符串转换为分类类型。

我们可能希望多次处理文件，例如尝试不同的查询。在这种情况下，我们不希望在加载后进行数据类型转换，而是希望数据以记住我们想要使用的数据类型的方式存储在磁盘上。CSV 不符合要求，因为它基本上只是一堆字符串。

此外，加载 CSV 可能会很慢，涉及大量解析。Parquet 数据格式 是一个更好的替代方案：它具有与 Pandas 相似的实际数据类型概念，并且加载速度更快。

作为第一步，我们将加载 CSV，选择更好的列类型，并将结果写入 Parquet 文件：

import sys
import pandas as pd

df = pd.read_csv(
    sys.argv[1],
    dtype={
        "route_id": "category",
        "direction_id": "category",
        "point_type": "category",
        "standard_type": "category",
    },
    parse_dates=["service_date", "scheduled", "actual"],
)

df.to_parquet(sys.argv[1].replace(".csv", ".parquet"))

额外的，Parquet 使用压缩：新文件大小为 20MB，而 CSV 为 300MB。请记住，这只是磁盘上的大小。数据在加载到内存之前必须解压缩，因此磁盘压缩对内存使用没有帮助。

第二步：使用简单的 Pandas 实现查找慢速公交路线

为了找到慢速公交路线，我们将专注于“headways”：特定公交路线的到达频率。如果我们查看上面的示例数据，我们可以看到入站公交 1 应该每 1200 秒到达一次，但在 5 月 1 日，它实际上到达得更快，相差 841 秒。并非所有行都有 headway 信息；我们只想要 standard_type 为 Headway 的行。

以下是我们的算法：

1. 删除所有没有 headway 信息的行。
2. 计算实际 headway 与预期 headway 的比率；如果大于 1，则意味着公交车晚点。
3. 对于每对路线编号和方向（入站/出站），选择该月的 headway 比率中位数。
4. 找到 headway 比率中位数最差的 5 对路线。

这可能不是找到慢速公交车的最佳方法，但我们只是将其用作示例，所以没关系。

以下是第一次尝试，一个简单的实现：

import pandas as pd

def find_worst_headways():
    # 加载数据：
    data = pd.read_parquet("MBTA-2022-05.parquet")
    # 过滤到仅包含 headway 信息的行：
    data = data[data["standard_type"] == "Headway"]
    # 计算实际 headway 与预期 headway 的比率：
    data["headway_ratio"] = (
        data["headway"] / data["scheduled_headway"]
    )
    # 按路线和方向（入站/出站）分组：
    by_route = data.groupby(["route_id", "direction_id"])
    # 找到每条路线的 headway 比率中位数：
    median_headway = by_route[["headway_ratio"]].median()
    # 返回最差的 5 条路线：
    return median_headway.nlargest(
        5, columns=["headway_ratio"]
    )

print(find_worst_headways())

以下是结果：

                       headway_ratio
route_id direction_id
108      Outbound           2.900000
88       Outbound           1.680000
83       Outbound           1.565000
134      Outbound           1.431111
         Inbound            1.346667

由于某种原因，一个路线 ID 是空白的。我还没有调查原因，因为这只是一个示例，但可能是输入数据格式不正确。只要我们后面的实现给出相同的结果，这对本文来说并不重要，我们只是在比较相同的东西。在 Polars 实现中，它显示为 134，其他结果相同。

通过使用 /usr/bin/time -v 运行程序，我们可以看到 最大 RSS（驻留）内存使用 和 挂钟时间和 CPU 时间：

User time (seconds): 0.84
System time (seconds): 1.33
Percent of CPU this job got: 491%
Elapsed (wall clock) time (h:mm:ss or m:ss): 0:00.44
Maximum resident set size (kbytes): 909500

Pandas 不是并行的，但 Parquet 加载库（在本例中为 Arrow）可以利用多个 CPU。

第三步：调查内存使用情况

我们已经了解到，我们的简单 Pandas 实现使用了 909MB 的内存。这很多！接下来，我们将使用 Sciagraph 性能和内存分析器 来测量内存使用的来源。

此报告声称分配了 1.2GB 的内存；之前我们看到最大驻留内存为 900MB。差异是由于 测量不同的东西。

以下是大部分内存分配的地方：

• 100MB 来自过滤到仅包含 headway 信息的行。
• 75MB 来自计算中位数。
• 1000MB 来自 Arrow 库，用于加载数据。这不是 Python 代码，Sciagraph 尚未显示本机调用堆栈的内存分配，因此尚不清楚加载数据的哪一部分负责。

显然，我们希望专注于最后一项，但我们在那里的详细信息较少。我们可以切换到 Memray 内存分析器，它确实提供了本机（C）调用堆栈。然而，稍微思考一下就会发现问题的一部分，以及明显的下一步。我们当前的处理涉及加载大量数据，然后丢弃大部分数据。

特别是，我们：

1. 加载所有数据；这是 Arrow 参与并分配大量内存的地方。
2. 删除许多行，特别是那些没有 headway 数据的行。
3. 忽略许多我们在此查询中未使用的数据列。

分块或批处理是 减少内存使用的基本技术之一。如果我们分块加载数据，而不是一次性加载所有数据，我们可以逐块过滤数据。然后我们可以合并更小的块，并在更少的数据上运行我们的逻辑。

第四步：更优化的 Pandas 实现

以下是我们基于新见解的实现：

import pandas as pd
import pyarrow.parquet as pq

def find_worst_headways():
    # 分块加载数据：
    chunks = []
    parquet_file = pq.ParquetFile(
        "MBTA-2022-05.parquet"
    )
    for batch in parquet_file.iter_batches():
        chunk = batch.to_pandas()
        del batch
        # 计算 headway 比率：
        chunk["headway_ratio"] = (
            chunk["headway"] / chunk["scheduled_headway"]
        )
        # 存储我们关心的列：
        chunks.append(chunk[
            ["route_id", "direction_id", "headway_ratio"]
        ])
    del parquet_file

    # 合并为一个大的 DataFrame。
    # 不理想，涉及两个内存副本...
    data = pd.concat(chunks)
    del chunks

    # 按路线和方向（入站/出站）分组：
    by_route = data.groupby(["route_id", "direction_id"])
    # 找到每条路线的 headway 比率中位数：
    median_headway = by_route[["headway_ratio"]].median()
    # 返回最差的 5 条路线：
    return median_headway.nlargest(
        5, columns=["headway_ratio"]
    )

print(find_worst_headways())

输出相同，但使用的内存更少：

User time (seconds): 1.10
System time (seconds): 1.26
Percent of CPU this job got: 348%
Elapsed (wall clock) time (h:mm:ss or m:ss): 0:00.68
Maximum resident set size (kbytes): 364172

我们已经从 900MB 的最大驻留内存减少到 360MB，这是一个非常好的改进。

第五步：尝试用 fastparquet 替换 PyArrow

在我们最初的内存分析中，我们看到 PyArrow 负责大部分分配的内存，作为加载 Parquet 文件的一部分。Pandas 还可以使用另一个名为 fastparquet 的库加载 Parquet 文件，因此我们可以将我们的简单版本和优化版本都切换到使用它，看看它如何影响内存使用。

以下是我们在简单版本中所做的更改：

# ...
data = pd.read_parquet("MBTA-2022-05.parquet",
                       engine="fastparquet")
# ...

优化版本：

import pandas as pd
import fastparquet as pq

def find_worst_headways():
    # 分块加载数据：
    chunks = []
    parquet_file = pq.ParquetFile("MBTA-2022-05.parquet")
    for chunk in parquet_file.iter_row_groups():
        # 计算 headway 比率：
        chunk["headway_ratio"] = (
            chunk["headway"] / chunk["scheduled_headway"]
        )
        # ...

在测量结果代码的内存使用情况时，事实证明 Fastparquet 在简单版本中使用的内存比 PyArrow 少得多。但优化版本实际上更糟！我可以花时间尝试找出原因，但这可能太偏离主题了。猜测一下，它正在加载整个文件，如果我们希望从分块中获得任何好处，我们需要在创建 Parquet 文件时调整行组的大小。

以下是我们各种实现的总结：

实现	最大驻留内存	挂钟时间	CPU 时间
Pandas 简单（PyArrow）	909MB	0.44 秒	2.17 秒
Pandas 优化（PyArrow）	364MB	0.68 秒	2.36 秒
Pandas 简单（Fastparquet）	400MB	0.71 秒	2.07 秒
Pandas “优化”（Fastparquet）	460MB	0.71 秒	2.04 秒

---

惰性处理，懒惰的程序员：使用 Polars 减少内存使用

到目前为止，我们已经了解到，手动实现的批处理实现至少在使用 PyArrow 时可以减少 Pandas 中的内存使用。令人恼火的是，这需要我们手动重构数据的表示和加载方式。理想情况下，我们的库会为我们做到这一点，但不幸的是，在使用 Pandas 时这是不可能的。

Pandas 是一个 急切的 API：你告诉它做某事，它会立即执行。 因此，如果你告诉它加载一个文件，它会立即将所有内容加载到内存中；它无法知道你打算在下一行代码中丢弃一半的数据。

另一种选择是 惰性 API，它允许你将一系列操作——加载、过滤、聚合、转换——串在一起，而无需实际执行任何工作。 在创建了这一系列操作之后，你可以单独告诉库执行整个操作。

一个聪明的惰性库可以查看所有操作，并制定一个优化的执行计划，考虑到你计划做的所有事情——以及你计划不做的事情。例如：

• 如果你根本不接触某一列，则无需将其加载到内存中。
• 如果可以进行批处理，库可以自动为你进行批处理。Polars 允许你明确要求这种优化；请参阅下面的流式处理。

Polars 是一个具有许多优点的 Pandas 替代品，例如多核处理——并且它支持急切和惰性 API。使用惰性 API 可以意味着减少内存使用，而无需额外的工作来手动批处理数据处理。

我们的 Polars 实现

Polars 的急切加载 API 通常以 read_* 开头，而惰性加载 API 以 scan_* 开头。以下是我们的代码在使用 Polars 惰性 API（特别是 scan_parquet()）重新实现时的样子：

import polars as pl

def headways_sorted_worst_first():
    # 惰性加载数据：
    data = pl.scan_parquet("MBTA-2022-05.parquet")
    # 过滤到仅包含 headway 信息的行，然后选择我们需要的数据：
    data = data.filter(
        pl.col("standard_type") == "Headway"
    ).select(
        [
            pl.col("route_id"),
            pl.col("direction_id"),
            pl.col("headway") / pl.col("scheduled_headway"),
        ]
    )
    # 按路线和方向（入站/出站）分组：
    by_route = data.groupby(["route_id", "direction_id"])
    # 找到每条路线的 headway 比率中位数：
    median_headway = by_route.agg(
        pl.col("headway").median()
    )
    # 没有 nlargest() 方法，因此只需按降序排序：
    return median_headway.sort("headway", reverse=True)

# 创建查询：
query = headways_sorted_worst_first()
# 实际运行查询：
result = query.collect()
# 打印最差的 5 个 headway：
print(result[:5, :])

当我们使用 Pandas 时，每次调用都会执行一些操作。使用 Polars 的惰性 API，在调用 collect() 之前实际上不会发生任何事情。此时，它可以使用查询计划器来制定优化的执行策略。只有在那时，Polars 才会执行加载、过滤和聚合数据的工作。

未来的 Polars 版本可能会改进查询计划器（或者，可能会变得更糟）。使用 Pandas 的急切 API，结构优化需要你重构代码；无论好坏，执行策略都取决于你。

内存使用和性能比较

让我们看看 Polars 的惰性实现与我们之前的实现相比如何：

实现	最大驻留内存	挂钟时间	CPU 时间
Pandas 简单（PyArrow）	909MB	0.44 秒	2.17 秒
Pandas 优化（PyArrow）	364MB	0.68 秒	2.36 秒
Pandas 简单（Fastparquet）	400MB	0.71 秒	2.07 秒
Pandas “优化”（Fastparquet）	460MB	0.71 秒	2.04 秒
Polars（惰性）	152MB	0.11 秒	0.44 秒

Polars 使用更少的内存，完成速度更快，并且使用的 CPU 资源也更少。这太棒了！更妙的是，我们不需要重新结构化代码来手动实现批处理；要么 Polars 为我们完成了这一步，要么它应用了其他一些减少内存使用的技术。真正的原因在于它的惰性 API：仅仅使用 Polars 的急切 read_parquet() API 加载文件，最大驻留内存就只有 310MB。

请注意，Polars 包含一个流模式（截至 2023 年 1 月仍处于实验阶段），它特别尝试使用批处理 API 来降低内存使用。只需调用 collect(streaming=True) 而不是 collect()。在这种情况下，它并没有显著减少内存使用，但如果你处理的是大文件，可以尝试一下。