从Pandas 到 Polars 六：在 Polars 中流式处理大型数据集

Polars 相对于 Pandas 的一大优势是，处理大于内存的数据集可以像向函数调用添加一个参数一样简单。然而，流式处理并不适用于所有情况。在这篇文章中，我将介绍流式处理是如何工作的，以及如何解决你可能面临的一些挑战。

在 Polars 中的流式处理

首先，什么是流式处理？在 Polars 中，当我们以批次方式处理数据而不是一次性处理所有数据时，我们称之为流式处理。基本上，这就像 Polars 对数据进行循环遍历，并依次处理每个批次。相比之下，当我们使用默认的一次性处理所有数据的方式时，我们在 Polars 中称之为标准引擎。

流式处理允许我们处理大于内存的数据集，因为 Polars 可以处理许多小批次的数据，而不是一个大型数据集。我们甚至不需要查询的最终输出完全适应内存，因为我们可以使用其中一个 sink_ 方法（见下文）将输出写入磁盘上的 Parquet 或其他文件。

在 Polars 中，只有当我们使用惰性模式时才会发生流式处理。我们必须明确地告诉 Polars 我们想要使用流式处理引擎，当评估一个惰性表达式时，我们需要传递 streaming=True（如下所示）。

并非所有惰性查询都可以在流式处理模式下执行。如果你尝试在流式处理模式下执行无法在流式处理模式下执行的查询，则 Polars 会默认使用标准引擎来运行查询。

实践中的流式处理

在这个惰性查询的示例中，我们：

• 对 CSV 文件进行惰性扫描
• 按列进行分组
• 计算另一列的平均值

我们通过向 collect 方法传递 streaming=True 来使用流式处理引擎执行这个惰性查询。

query = pl.scan_csv("iris.csv").group_by("species").agg(

    pl.col("sepal_width").mean().alias("mean_sepal_width")

)

query.collect(streaming=True)

另外，当我们传递 streaming=True 时，Polars 不保证查询会使用流式处理引擎执行。因此，始终检查你的查询是否实际使用了流式处理引擎是一个好习惯。

我们可以通过使用 explain 方法打印优化后的查询计划来进行此检查。流式处理引擎的优化查询计划与标准引擎的优化查询计划是不同的，所以我们也需要向 explain 传递 streaming=True。

query.explain(streaming=True)

这给我们以下输出：

--- STREAMING

AGGREGATE

[col("sepal_width").mean().alias("mean_sepal_width")] BY [col("species")] FROM


    Csv SCAN iris.csv

    PROJECT 2/5 COLUMNS  --- END STREAMING


  DF []; PROJECT */0 COLUMNS; SELECTION: "None"

我们看到整个查询都在以“--- STREAMING”开始并以“--- END STREAMING”结束的区块中。这告诉我们整个查询正在以流式模式运行。

流式处理引擎将查询分成批次。批次的大小取决于您机器上的 CPU 数量以及查询中每行所需的内存量。后面更详细地探讨这一点。

直接流式写入文件

如果查询的输出仍然太大而无法装入内存，我们可以直接将输出写入磁盘上的文件。这可以通过使用 sink_ 方法之一来实现。在这个例子中，我们将输出写入一个 Parquet 文件。

query.sink_parquet("iris.parquet")

我们不需要在这里传递 streaming=True，因为 sink_parquet 仅使用流式处理引擎运行。如果查询的某一部分不在流式模式下运行，sink_parquet 会失败并显示错误消息。

为什么流式处理不适用于所有查询？

流式处理并不适用于所有查询。在这个例子中，我们对数据集中的某一列进行了累积求和（cum_sum）操作，如下所示：

query = (

    pl.scan_csv("iris.csv")

    .with_columns(

        pl.col("sepal_width").cum_sum().alias("cum_sum_sepal_width")

    ))

当我们使用 explain(streaming=True) 打印查询计划时，我们得到以下输出：

 WITH_COLUMNS:

 [col("sepal_width").cum_sum().alias("cum_sum_sepal_width")]

  --- STREAMING

  Csv SCAN iris.csv

  PROJECT */5 COLUMNS  --- END STREAMING

    DF []; PROJECT */0 COLUMNS; SELECTION: "None"

我们看到，查询计划中与 cum_sum 相关的 WITH_COLUMNS 部分并不在流式处理块中。

那么为什么这个查询不能在流式模式下运行呢？要理解这一点，请注意 Polars 的开发者需要在流式引擎中实现每个操作，这些操作可能是一个方法（如 group_by）或一个表达式（如 .mean）。作为一个一般原则，能够一次应用于一批数据然后将结果合并的操作更适合流式引擎，而不是那些一批数据需要其他批次信息的操作。

例如，如果我们计算一列的总和，我们可以分批次进行，先计算每个批次的和，然后再将所有批次的结果相加。然而，如果我们想要计算累积的 cum_sum，则需要先对第一批数据的值进行累积求和，然后将最终值传递给第二批数据，依此类推。这使得 cum_sum 更难被包含在流式引擎中。

流式查询的故障排查

如果你在流式模式下执行查询时遇到困难，我建议你按照以下步骤操作：

• 检查查询计划，查看查询的哪些部分可以在流式模式下执行
• 如果查询计划中有部分无法在流式模式下执行，并且你不确定是哪个操作导致的，可以尝试从一个更简单的查询开始，并逐渐增加部分，直到找到导致问题的操作
• 如果你找到一个操作意味着查询无法在流式模式下执行，尝试将该操作推迟到查询中数据大小已经减少的某个点，例如在 filter 或 group_by 之后

在某些情况下，对于在流式模式下无法工作的查询，有替代解决方案。例如，假设我们有一个包含时间序列数据的 DataFrame

from datetime import datetime
import polars as pl

df = (
    pl.DataFrame(
        {
"datetime":pl.datetime_range(datetime(2024,1,1),datetime(2024,1,3),"1h",eager=True)
        }
    )
.with_row_index()
)

shape: (49, 2)

我们想要对数据进行按日分组。通常我们会使用 group_by_dynamic 方法来进行时间分组。然而，这个方法目前并不支持流式模式。

不过，我们仍然可以在流式模式下通过 with_columns 来从日期时间中提取日期，然后在日期列上进行常规的 group_by 操作来实现按日分组。

groupby_query = (
    df
    .lazy()
    .with_columns(pl.col("datetime").dt.date().alias("date"))
    .group_by("date")
    .agg(pl.col("index").mean()))

如果我们打印流式查询计划，我们会得到以下输出：

--- STREAMING
AGGREGATE
[col("index").mean()] BY [col("date")] FROM
   WITH_COLUMNS:
   [col("datetime").dt.date().alias("date")]
    DF ["index", "datetime"]; PROJECT 2/2 COLUMNS; SELECTION: "None"  --- END STREAMING

  DF []; PROJECT */0 COLUMNS; SELECTION: "None"

在这里，我们看到查询的 WITH_COLUMNS 和 AGGREGATE 部分都在流式处理块中，因此这个完整的查询将在流式模式下运行。

你可能能够识别出对你的查询进行类似转换的方法，以便让它在流式模式下运行。

以上就是Polars中流式模式工作核心思想的简单介绍。