从 Pandas 到 Polars 十三：流式处理的关键参数

在本文中，我们将看到Polars如何设置一些影响流式模式的关键参数。如果你想要优化大型流式查询的性能或处理流式查询中的内存不足错误，理解这些概念是很重要的。

在流式模式下，Polars会分块处理数据，而不是一次性处理所有数据。这使得Polars能够处理大于内存的数据集，而无需你付出太多额外的努力。这就引出了一个问题——Polars是如何决定块的大小的？

当流式模式首次被引入时，答案很简单，因为块的大小被设置为50000行。最近，块的大小已经更改为DataFrame中列的数量和可用线程数量的函数。

创建一个简单的示例数据集

我们将为这个示例创建一个简单的测试集，包括3个带有随机数据的CSV文件：

import polars as plimport numpy as np

for i in range(3):
    (
        pl.DataFrame(
            np.random.standard_normal((10,3))
        )
        .with_row_count('id')
        .write_csv(f"test_{i}.csv")
    )

每个DataFrame看起来像这样：

shape: (10, 4)
┌─────┬───────────┬───────────┬───────────┐
│ id  ┆ column_0  ┆ column_1  ┆ column_2  │
│ --- ┆ ---       ┆ ---       ┆ ---       │
│ i64 ┆ f64       ┆ f64       ┆ f64       │
╞═════╪═══════════╪═══════════╪═══════════╡
│ 0   ┆ 1.189349  ┆ -0.288161 ┆ -0.156299 │
│ 1   ┆ 0.602317  ┆ -0.108463 ┆ -1.029294 │
│ 2   ┆ 0.230017  ┆ 0.18799   ┆ 1.251513  │
│ 3   ┆ -0.706367 ┆ 0.407965  ┆ -0.509364 │
│ …   ┆ …         ┆ …         ┆ …         │
│ 6   ┆ -0.996497 ┆ 0.359672  ┆ -0.382435 │
│ 7   ┆ 0.310096  ┆ -1.165665 ┆ 1.444398  │
│ 8   ┆ -1.427343 ┆ -0.762879 ┆ -0.077449 │
│ 9   ┆ 1.248644  ┆ 0.498989  ┆ 0.96604   │
└─────┴───────────┴───────────┴───────────┘

我们使用一个简单的查询，其中我们按id列进行分组，并对浮点数字段进行求和：

(
    pl.scan_csv("test_*.csv")
    .groupby('id')
    .agg(
        pl.col(pl.Float64).sum()
    ))

当我们使用pl.scan_csv时，这是一个惰性查询，所以此时不会发生任何事情。

对于任何惰性查询，我们都可以使用explain方法来查看优化后的查询计划。对于这个查询，如果不对explain方法传递任何参数，我们会看到：

 print(
        (
            pl.scan_csv("test_*.csv")
            .groupby('id')
            .agg(
                pl.col(pl.Float64).sum()
            )
            .explain()
        )
 )
 AGGREGATE
[col("column_0").sum(), col("column_1").sum(), col("column_2").sum()] BY [col("id")] FROM
FAST_PROJECT: [column_0, column_1, column_2, id]
  UNION:
  PLAN 0:

      CSV SCAN test_0.csv
      PROJECT */4 COLUMNS
  PLAN 1:

      CSV SCAN test_1.csv
      PROJECT */4 COLUMNS
  PLAN 2:

      CSV SCAN test_2.csv
      PROJECT */4 COLUMNS
  END UNION

我们从下往上读取这个查询计划。

1. CSV SCAN节点读取CSV文件，而PROJECT部分选择我们想要的列（在这种情况下是所有4列）。
2. UNION节点将来自三个CSV文件的结果合并成一个单一的DataFrame。
3. AGGREGATE节点执行分组和求和操作。

这是如果我们使用默认的非流式引擎和.collect()方法时，Polars将执行的查询计划。

我的查询是在流式模式下运行的吗？

如果我尝试使用.collect(streaming=True)在流式模式下执行查询，会发生什么？我们可以通过向explain方法传递streaming=True来查看这种模式的查询计划：

print(
    pl.scan_csv("test*.csv")
    .groupby("id")
    .agg(
        pl.col(pl.Float64).sum()
    )
    .explain(streaming=True))STREAMING CHUNK SIZE: 12500 rows--- PIPELINEAGGREGATE
[col("column_0").sum(), col("column_1").sum(), col("column_2").sum()] BY [col("id")] FROM
FAST_PROJECT: [column_0, column_1, column_2, id]
  UNION:
  PLAN 0:

      CSV SCAN test_0.csv
      PROJECT */4 COLUMNS
  PLAN 1:

      CSV SCAN test_1.csv
      PROJECT */4 COLUMNS
  PLAN 2:

      CSV SCAN test_2.csv
      PROJECT */4 COLUMNS
  END UNION  --- END PIPELINE

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

--- PIPELINE 和 --- END PIPELINE 这两行告诉我们查询将在流式模式下执行。而 STREAMING CHUNK SIZE 告诉我们 Polars 将以（最多）12500 行的数据块来处理数据。

请注意 - 这个 explain(streaming=True) 功能目前还处于实验阶段，并且是 alpha 阶段。

如何计算数据块大小？

在流式模式下影响性能的一个关键参数是数据块大小。数据块大小是 Polars 在每个批次中处理的行数。在流式查询计划的顶部，我们看到数据块大小是 12500 行。这个值是由 Polars 根据以下因素设置的：

1. Polars 将使用的线程数
2. 数据集中的列数

我们将在下面看到如何修改数据块大小，但首先我们将查看 Polars 如何为每个查询计算默认的数据块大小。

数据块大小的公式是列数和线程数的函数。我们首先看线程数。

线程数是多少？

在查看数据块大小公式之前，我们需要了解 Polars 将使用多少线程。默认情况下，线程数设置为你机器上的 CPU 数量（或者 Docker 容器内可用的最大 CPU 数量）。

你可以使用 pl.threadpool_size() 函数来检查 Polars 将使用的最大线程数。

import polars as plpl.threadpool_size()

在我的电脑上，这个值通常是12；但如果我在 Docker 容器里运行，并且我把最大 CPU 数量设置为7，那么它就会是7

你可以通过设置POLARS_MAX_THREADS环境变量来覆盖线程的最大数量。

import osos.environ["POLARS_MAX_THREADS"] = "4"import polars as pl

请注意 - 你必须在导入Polars之前设置这个环境变量。如果你在导入Polars之后再设置它，将不会产生任何效果。

就我而言，我发现默认的线程数工作得很好，我几乎从不改变它。然而，如果你在使用Polars的同时运行其他CPU密集型任务，你可能想要减少线程数。

如果你有一台专门用于查询的机器，你可能想要增加线程数。我对此的实验结果并不明确——在某些情况下，我发现增加线程数可以提高性能，而在其他情况下，它几乎没有效果或者会降低性能。

分块大小公式

设置分块大小的当前公式简而言之就是：

        thread_factor = max(12 / n_threads, 1)
        STREAMING_CHUNK_SIZE = max(50_000 / n_cols * thread_factor, 1000))

thread_factor 的设置最大为 1，且为 12/n_threads 的结果。如果你有 12 个或更多的线程，thread_factor 将为 1；如果你的线程数少于 12，thread_factor 将大于 1。

分块大小（chunk size）的设置是以下两者中的较大值：

1. 1000 或
2. 50000 除以 DataFrame 中的列数，再乘以 n_threads。

在我有 12 个线程的情况下，thread_factor 是 1，所以分块大小就是 50,000 除以 LazyFrame 中的列数。由于我的 LazyFrame 有 4 列，所以分块大小是 12,500 行。

我预计这个公式会随着 Polars 针对不同用例的优化而不断演变。

手动设置分块大小

分块大小的最优值将取决于您的数据、查询和硬件。您可以使用 pl.Config.set_streaming_chunk_size 函数手动设置分块大小。在这个例子中，我将它设置为 50,000 行：

pl.Config.set_streaming_chunk_size(50000)

与 POLARS_MAX_THREADS 环境变量不同，你可以在任何时候设置分块大小。例如，如果我将分块大小设置为 50,000 行，然后再次运行 explain(streaming=True)，我会得到以下输出：

STREAMING CHUNK SIZE: 50000 rows--- PIPELINE
AGGREGATE
[col("column_0").sum(), col("column_1").sum(), col("column_2").sum()] BY [col("id")] FROM
FAST_PROJECT: [column_0, column_1, column_2, id]
  UNION:
  PLAN 0:

      CSV SCAN test_0.csv
      PROJECT */4 COLUMNS
  PLAN 1:

      CSV SCAN test_1.csv
      PROJECT */4 COLUMNS
  PLAN 2:

      CSV SCAN test_2.csv
      PROJECT */4 COLUMNS
  END UNION  --- END PIPELINE

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

这里我们可以看到流式处理分块大小现在是 50,000 行。

如果你正在处理大型数据集，建议你尝试不同的分块大小，以找到适用于你的数据和硬件的最优值——这可能会对性能产生重大影响。