今天我們來了解Polars的lazy模式。

藉由lazy模式，Polars能事先利用各種最佳化技巧打造最佳的query plan，大幅提升相對於eager模式的計算效率。

本日大綱如下：

0. 本日引入模組及準備工作
1. 生成pl.LazyFrame
2. 執行pl.LazyFrame所記錄的運算（materialize）
3. Query plan
4. pl.collect_all()
5. pl.LazyFrame的使用限制
6. 實際應用
7. codepanda

0. 本日引入模組及準備工作

import io
import polars as pl

data = {"col1": [1, 2, 3], "col2": ["x", "y", "z"]}
df = pl.DataFrame(data)


def get_csv_text():
    csv_text = "col1,col2\n1,x\n2,y\n3,z\n"
    return io.StringIO(csv_text)

為方便講解，我們使用io.StringIO來做為後續pl.scan_*()的輸入。

可以假想為其是一個含有以下內容的csv檔案：

col1,col2
1,x
2,y
3,z

1. 生成pl.LazyFrame

pl.LazyFrame可以想做是一張草稿，上面記錄了所有運算過程及相關資訊，有以下兩種生成方式。

使用pl.scan_*()

相比於eager模式的pl.read_*()函數，lazy模式的函數則為pl.scan_*()。舉例來說，使用pl.scan_csv()可以讀取csv檔案並返回pl.LazyFrame。

pl.scan_csv(get_csv_text())

naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)

Csv SCAN [22 in-mem bytes]
PROJECT */2 COLUMNS

呼叫pl.DataFrame.lazy()

使用pl.DataFrame.lazy()可以將一個pl.DataFrame轉換為pl.LazyFrame()：

df.lazy()

naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)

DF ["col1", "col2"]; PROJECT */2 COLUMNS

2. 執行pl.LazyFrame所記錄的運算（materialize）

API文件中提到，只有在呼叫以下四種函數時，pl.LazyFrame才會正式執行其所記錄的運算（Polars稱呼此動作為materialize）：

• pl.LazyFrame.collect()：執行運算，並返回pl.DataFrame。
• pl.LazyFrame.explain()：印出query plan。
• pl.LazyFrame.show_graph()：使用graphviz graph印出query plan。
• pl.LazyFrame.collect_schema()：返回最終schema。

pl.LazyFrame.collect()是最常使用的功能，可以實際執行pl.LazyFrame所記錄的運算，並返回pl.DataFrame：

pl.scan_csv(get_csv_text()).collect()

shape: (3, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ str  │
╞══════╪══════╡
│ 1    ┆ x    │
│ 2    ┆ y    │
│ 3    ┆ z    │
└──────┴──────┘

3. Query plan

使用pl.LazyFrame.explain()或pl.LazyFrame.show_graph()可以印出Polars在materialize所會實際運行的計算。

舉例來說，在q1中我們需要進行將「"col2"」列轉換為大寫及選取「"col1"」列中小於或等於2的行數：

q1 = (
    pl.scan_csv(get_csv_text())
    .with_columns(pl.col("col2").str.to_uppercase())
    .filter(pl.col("col1").le(2))
)
q1.explain()

WITH_COLUMNS:
 [col("col2").str.uppercase()] 
  Csv SCAN [22 in-mem bytes]
  PROJECT */2 COLUMNS
  SELECTION: [(col("col1")) <= (2)]

觀察query plan得知（註1），pl.scan_csv()不必將所有資訊都先讀進記憶體，可以在讀取前就先判斷最終所需要的列及行，例如q1中的「"col1"」的最後一行為3將會被篩掉，所以根本不用讀進來也不用進行轉換大小寫的工作。

materialize後可得到pl.DataFrame：

q1.collect()

shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ str  │
╞══════╪══════╡
│ 1    ┆ X    │
│ 2    ┆ Y    │
└──────┴──────┘

教學文件中有列出數種Polars常用的最佳化技巧，上面這種針對pl.DataFrame.filter()最佳化的技巧稱之為「"predicate pushdown"」。

Polars預設會開啟所有能最佳化的部份，如果想得到沒有最佳化的query plan，可以將pl.LazyFrame.explain()的optimized=設為False：

q1.explain(optimized=False)

FILTER [(col("col1")) <= (2)]
FROM
   WITH_COLUMNS:
   [col("col2").str.uppercase()] 
    Csv SCAN [22 in-mem bytes]
    PROJECT */2 COLUMNS

此外，型別確認也可以視為是最佳化的一環。舉例來說，如果是想將「"col1"」列轉變為大寫：

❌
(
    pl.scan_csv(get_csv_text())
    .with_columns(pl.col("col1").str.to_uppercase())
    .collect()
)
# SchemaError: invalid series dtype: expected `String`, 
# got `i64` for series with name `col1`

但是由於「"col1"」列為pl.Int64型別，並沒有str命名空間可以使用，故而報錯。

請注意，此處並非是先讀取內容，直到materialize階段才因為計算失敗報錯，而是在materialize前，就先對比了schema，發現型別不匹配而提前報錯。

4. pl.collect_all()

當一個pl.LazyFrame想要進行不同的運算時，可以不必分開materialize，僅需使用一次pl.collect_all()（註2），這將使得Polars有機會可以查看各自的query plan，找出其中可以重覆使用的部份，進而提升效率。舉例來說，如果想針對lf1進行pl.LazyFrame.filter()及pl.LazyFrame.select()，可以這麼寫：

lf1 = pl.scan_csv(get_csv_text())
lf1_a = lf1.filter(pl.col("col1").le(2))
lf1_b = lf1.select(pl.col("col2"))
pl.collect_all([lf1_a, lf1_b])

[shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ str  │
╞══════╪══════╡
│ 1    ┆ x    │
│ 2    ┆ y    │
└──────┴──────┘, shape: (3, 1)
┌──────┐
│ col2 │
│ ---  │
│ str  │
╞══════╡
│ x    │
│ y    │
│ z    │
└──────┘]

其返回值，為一含有多個pl.DataFrame的列表，代表各個query的結果。

5. pl.LazyFrame的使用限制

由於lazy模式中，需要確切掌握所有型別，才能打造最佳的query plan，所以部份pl.DataFrame所提供的功能，pl.LazyFrame並沒有提供，例如pivot。

教學文件中提到的建議作法是在必要的操作前轉換為pl.DataFrame，並於操作完後轉換回pl.LazyFrame。例如：

(
    pl.LazyFrame({"id": ["a", "b", "c"], "month": ["jan", "feb", "mar"]})
    .with_columns(values=pl.Series([0, 1, 2]))
    .collect()
    .pivot(
        index="id", on="month", values="values", aggregate_function="first"
    )
    .lazy()
    .filter(pl.col("jan").is_null())
    .collect()
)

shape: (2, 4)
┌─────┬──────┬──────┬──────┐
│ id  ┆ jan  ┆ feb  ┆ mar  │
│ --- ┆ ---  ┆ ---  ┆ ---  │
│ str ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╡
│ b   ┆ null ┆ 1    ┆ null │
│ c   ┆ null ┆ null ┆ 2    │
└─────┴──────┴──────┴──────┘

此例中，由於含有pivot運算，所以我們先進行pl.LazyFrame.collect()得到pl.DataFrame後才進行pivot運算。運算完成後，馬上使用pl.DataFrame.lazy()轉換為回pl.LazyFrame進行後續運算。

6. 實際應用

舉一個常用的情境為例，當第一次處理很大的csv檔案時，如果只想先看看第一行，可能會直覺地寫下：

pl.read_csv(get_csv_text()).head(1)

此時，Polars會讀取整個csv後，顯示第一行。但如果是改用lazy寫法的話，Polars將只會讀取第一行：

✔️ # better
pl.scan_csv(get_csv_text()).head(1).collect()

7. codepanda

原生的Pandas為eager模式，但藉由dask的輔助，將能夠依靠lazy模式提升效率。

備註

註1：query plan的閱讀方法為由下往上，一段一段地閱讀，可以參考教學文件的說明。

註2：當使用pl.collect_all()時，有時候會遇到一種tricky的情形，有興趣的讀者可以自行閱讀教學文件。

Code

本日程式碼傳送門。