跳到内容

查询执行

我们在 Reddit 数据集上的示例查询是

scan_csv

q1 = (
    pl.scan_csv("docs/assets/data/reddit.csv")
    .with_columns(pl.col("name").str.to_uppercase())
    .filter(pl.col("comment_karma") > 0)
)

如果我们在 Reddit CSV 上运行上述代码,查询将不会被评估。相反,Polars 会获取每一行代码,将其添加到内部查询图并优化查询图。

当我们执行代码时,Polars 默认会执行优化后的查询图。

在完整数据集上执行

我们可以通过在查询上调用 .collect 方法,在完整数据集上执行查询。

scan_csv · collect

q4 = (
    pl.scan_csv(f"docs/assets/data/reddit.csv")
    .with_columns(pl.col("name").str.to_uppercase())
    .filter(pl.col("comment_karma") > 0)
    .collect()
)

shape: (14_029, 6)
┌─────────┬───────────────────────────┬─────────────┬────────────┬───────────────┬────────────┐
│ id      ┆ name                      ┆ created_utc ┆ updated_on ┆ comment_karma ┆ link_karma │
│ ---     ┆ ---                       ┆ ---         ┆ ---        ┆ ---           ┆ ---        │
│ i64     ┆ str                       ┆ i64         ┆ i64        ┆ i64           ┆ i64        │
╞═════════╪═══════════════════════════╪═════════════╪════════════╪═══════════════╪════════════╡
│ 6       ┆ TAOJIANLONG_JASONBROKEN   ┆ 1397113510  ┆ 1536527864 ┆ 4             ┆ 0          │
│ 17      ┆ SSAIG_JASONBROKEN         ┆ 1397113544  ┆ 1536527864 ┆ 1             ┆ 0          │
│ 19      ┆ FDBVFDSSDGFDS_JASONBROKEN ┆ 1397113552  ┆ 1536527864 ┆ 3             ┆ 0          │
│ 37      ┆ IHATEWHOWEARE_JASONBROKEN ┆ 1397113636  ┆ 1536527864 ┆ 61            ┆ 0          │
│ …       ┆ …                         ┆ …           ┆ …          ┆ …             ┆ …          │
│ 1229384 ┆ DSFOX                     ┆ 1163177415  ┆ 1536497412 ┆ 44411         ┆ 7917       │
│ 1229459 ┆ NEOCARTY                  ┆ 1163177859  ┆ 1536533090 ┆ 40            ┆ 0          │
│ 1229587 ┆ TEHSMA                    ┆ 1163178847  ┆ 1536497412 ┆ 14794         ┆ 5707       │
│ 1229621 ┆ JEREMYLOW                 ┆ 1163179075  ┆ 1536497412 ┆ 411           ┆ 1063       │
└─────────┴───────────────────────────┴─────────────┴────────────┴───────────────┴────────────┘

上面我们看到,在 1000 万行数据中,有 14,029 行符合我们的谓词。

使用默认的 collect 方法,Polars 会将所有数据作为一个批次处理。这意味着在查询内存使用峰值时,所有数据都必须适合您的可用内存。

复用 LazyFrame 对象

请记住,LazyFrame 是查询计划,即对计算的承诺,并且不保证缓存常见的子计划。这意味着每次在定义后将其在单独的下游查询中复用时,它都会重新计算。如果在 LazyFrame 上定义的操作不保留行顺序(例如 group_by),则每次运行时顺序也会改变。为避免这种情况,请对此类操作使用 maintain_order=True 参数。

在大于内存的数据上执行

如果您的数据所需的内存超过可用内存,Polars 可能会使用流式模式分批处理数据。要使用流式模式,只需将 engine="streaming" 参数传递给 collect

scan_csv · collect

q5 = (
    pl.scan_csv(f"docs/assets/data/reddit.csv")
    .with_columns(pl.col("name").str.to_uppercase())
    .filter(pl.col("comment_karma") > 0)
    .collect(engine='streaming')
)

在部分数据集上执行

当您在大型数据集上编写、优化或检查查询时,查询所有可用数据可能会导致开发过程变慢。

相反,您可以扫描分区的一个子集,或分别在查询的开头和结尾使用 .head/.collect。请记住,对数据子集进行聚合和过滤的结果可能无法代表您在完整数据上获得的结果。

scan_csv · collect · head

q9 = (
    pl.scan_csv(f"docs/assets/data/reddit.csv")
    .head(10)
    .with_columns(pl.col("name").str.to_uppercase())
    .filter(pl.col("comment_karma") > 0)
    .collect()
)

shape: (1, 6)
┌─────┬─────────────────────────┬─────────────┬────────────┬───────────────┬────────────┐
│ id  ┆ name                    ┆ created_utc ┆ updated_on ┆ comment_karma ┆ link_karma │
│ --- ┆ ---                     ┆ ---         ┆ ---        ┆ ---           ┆ ---        │
│ i64 ┆ str                     ┆ i64         ┆ i64        ┆ i64           ┆ i64        │
╞═════╪═════════════════════════╪═════════════╪════════════╪═══════════════╪════════════╡
│ 6   ┆ TAOJIANLONG_JASONBROKEN ┆ 1397113510  ┆ 1536527864 ┆ 4             ┆ 0          │
└─────┴─────────────────────────┴─────────────┴────────────┴───────────────┴────────────┘

分歧查询

查询在某个点发散是很常见的。在这种情况下,建议使用 collect_all,因为它们将确保发散查询只执行一次。

# Some expensive LazyFrame
lf: LazyFrame

lf_1 = LazyFrame.select(pl.all().sum())

lf_2 = lf.some_other_computation()

pl.collect_all([lf_1, lf_2]) # this will execute lf only once!