今天我們來說明Polars的context與expression概念（註1）。

context是Polars表達操作意圖的稱呼，可以分為三類操作：

• 欄：新增或修改欄位（DataFrame.with_columns()及DataFrame.select()）。
• 行：選取行數（DataFrame.filter()）。
• 分組聚合：依據某些條件進行分組後，以統計值或資料結構來代表各個分組（DataFrame.group_by().agg()）。

expression是Polars表達具體操作的稱呼，其種類繁多，幾乎所有您想做的操作都能夠用expression來完成。expression可以獨立定義在context之外，但大部份必須在context範圍內才會發揮作用。您可以將expression想成一種函數的語法糖，其可以在需要時，自動補抓到所對應的DataFrame及所在的context，而不需要顯性傳入。

由於Polars內建有query最佳化引擎，其會盡量減少各context內expression所需的計算。此外，各context內的所有expression都可以進行平行運算（Polars是以Rust寫成），所以執行速度相當快。最後，Polars的外掛系統允許您將Rust編寫的函數，註冊至Polars當作expression使用，所以即使有目前未支援的操作，使用者也有機會自己補上這個缺口。

接下來我們將介紹如何在三種context情境下使用expression（註2）。

以下這個名為df的Polars Dataframe（註3）將作為今天的範例：

import numpy as np
import polars as pl


np.random.seed(42)
data = {"nrs": [1, 2, 3, 4, 5], "random": np.random.rand(5)}
df = pl.DataFrame(data)

shape: (5, 2)
┌─────┬──────────┐
│ nrs ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 1   ┆ 0.37454  │
│ 2   ┆ 0.950714 │
│ 3   ┆ 0.731994 │
│ 4   ┆ 0.598658 │
│ 5   ┆ 0.156019 │
└─────┴──────────┘

於文章的最後，我將略為提及如何使用[]來存取行或欄。

欄

DataFrame.with_columns()

DataFrame.with_columns()可以新增欄位。

舉例來說，如果想新增一欄「"a"」欄位，其值分別為6、7、8、9、10五個整數，在沒接觸過Polars前，您可能會這麼寫：

❌
(df.with_columns(a=[6, 7, 8, 9, 10]))

shape: (5, 3)
┌─────┬──────────┬──────────────┐
│ nrs ┆ random   ┆ a            │
│ --- ┆ ---      ┆ ---          │
│ i64 ┆ f64      ┆ list[i64]    │
╞═════╪══════════╪══════════════╡
│ 1   ┆ 0.37454  ┆ [6, 7, … 10] │
│ 2   ┆ 0.950714 ┆ [6, 7, … 10] │
│ 3   ┆ 0.731994 ┆ [6, 7, … 10] │
│ 4   ┆ 0.598658 ┆ [6, 7, … 10] │
│ 5   ┆ 0.156019 ┆ [6, 7, … 10] │
└─────┴──────────┴──────────────┘

因為Polars是向量化思維，a=[6, 7, 8, 9, 10]會被解讀為想要「"a"」中的每一行都設定為[6, 7, 8, 9, 10]。

正確的寫法應該是使用像是pl.arange()的expression：

(df.with_columns(a=pl.arange(6, 11)))

shape: (5, 3)
┌─────┬──────────┬─────┐
│ nrs ┆ random   ┆ a   │
│ --- ┆ ---      ┆ --- │
│ i64 ┆ f64      ┆ i64 │
╞═════╪══════════╪═════╡
│ 1   ┆ 0.37454  ┆ 6   │
│ 2   ┆ 0.950714 ┆ 7   │
│ 3   ┆ 0.731994 ┆ 8   │
│ 4   ┆ 0.598658 ┆ 9   │
│ 5   ┆ 0.156019 ┆ 10  │
└─────┴──────────┴─────┘

如果欄位名稱不符合Python命名原則的話，可以使用Expr.alias()，像是：

(df.with_columns(pl.arange(6, 11).alias("*a*")))

shape: (5, 3)
┌─────┬──────────┬─────┐
│ nrs ┆ random   ┆ *a* │
│ --- ┆ ---      ┆ --- │
│ i64 ┆ f64      ┆ i64 │
╞═════╪══════════╪═════╡
│ 1   ┆ 0.37454  ┆ 6   │
│ 2   ┆ 0.950714 ┆ 7   │
│ 3   ┆ 0.731994 ┆ 8   │
│ 4   ┆ 0.598658 ┆ 9   │
│ 5   ┆ 0.156019 ┆ 10  │
└─────┴──────────┴─────┘

如果新增欄位需要用到既有欄位資訊的話，例如想新增一欄「"a"」欄位，其值為「"nrs"」欄位加1，可以這麼寫：

(df.with_columns(a=pl.col("nrs").add(1)))

shape: (5, 3)
┌─────┬──────────┬─────┐
│ nrs ┆ random   ┆ a   │
│ --- ┆ ---      ┆ --- │
│ i64 ┆ f64      ┆ i64 │
╞═════╪══════════╪═════╡
│ 1   ┆ 0.37454  ┆ 2   │
│ 2   ┆ 0.950714 ┆ 3   │
│ 3   ┆ 0.731994 ┆ 4   │
│ 4   ┆ 0.598658 ┆ 5   │
│ 5   ┆ 0.156019 ┆ 6   │
└─────┴──────────┴─────┘

這裡我們使用pl.col()來選取欄位，這將使得我們於取得該欄後可以使用expression提供的操作，像是這裡的Expr.add()。

由於context是平行運算，所以在同一個context中靠後面的運算，無法參考前方的計算結果。舉例來說，如果想新增一欄「"a"」欄位，其值為「"nrs"」欄位加1，且想再新增一欄「"b"」欄位，其值為「"a"」欄位加1，必須這麼寫：

(df.with_columns(pl.col("nrs").alias("a")).with_columns(pl.col("a").add(1).alias("b")))

shape: (5, 4)
┌─────┬──────────┬─────┬─────┐
│ nrs ┆ random   ┆ a   ┆ b   │
│ --- ┆ ---      ┆ --- ┆ --- │
│ i64 ┆ f64      ┆ i64 ┆ i64 │
╞═════╪══════════╪═════╪═════╡
│ 1   ┆ 0.37454  ┆ 1   ┆ 2   │
│ 2   ┆ 0.950714 ┆ 2   ┆ 3   │
│ 3   ┆ 0.731994 ┆ 3   ┆ 4   │
│ 4   ┆ 0.598658 ┆ 4   ┆ 5   │
│ 5   ┆ 0.156019 ┆ 5   ┆ 6   │
└─────┴──────────┴─────┴─────┘

如果寫在同一個context中的話，將會引發ColumnNotFoundError。

❌
(df.with_columns(pl.col("nrs").alias("a"), pl.col("a").add(1).alias("b")))

ColumnNotFoundError: a

此外，DataFrame.with_columns()也可以用來修改欄位，例如將「"nrs"」欄位每行都減1：

(df.with_columns(pl.col("nrs").sub(1)))

shape: (5, 2)
┌─────┬──────────┐
│ nrs ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 0   ┆ 0.37454  │
│ 1   ┆ 0.950714 │
│ 2   ┆ 0.731994 │
│ 3   ┆ 0.598658 │
│ 4   ┆ 0.156019 │
└─────┴──────────┘

甚至直接將整欄換成另一個型別，例如：

(df.with_columns(nrs=pl.Series(list("abcde")).explode()))

shape: (5, 2)
┌─────┬──────────┐
│ nrs ┆ random   │
│ --- ┆ ---      │
│ str ┆ f64      │
╞═════╪══════════╡
│ a   ┆ 0.37454  │
│ b   ┆ 0.950714 │
│ c   ┆ 0.731994 │
│ d   ┆ 0.598658 │
│ e   ┆ 0.156019 │
└─────┴──────────┘

這裡我使用Series.explode()來將list("abcde")分配到「"nrs"」欄位的每一行。

DataFrame.select()

DataFrame.select()可以讓我們選取想要的欄位，例如：

(df.select("nrs"))
(df.select(pl.col("nrs"))

shape: (5, 1)
┌─────┐
│ nrs │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 2   │
│ 3   │
│ 4   │
│ 5   │
└─────┘

可以直接使用欄位名或透過pl.col()兩種方式來選擇。

如果選取欄位不存在的話，則會當場生成，例如：

(df.select("nrs", pl.col("nrs").alias("a")))

shape: (5, 2)
┌─────┬─────┐
│ nrs ┆ a   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 2   ┆ 2   │
│ 3   ┆ 3   │
│ 4   ┆ 4   │
│ 5   ┆ 5   │
└─────┴─────┘

此處我們選擇「"nrs"」欄位並當場生成「"a"」欄位。

此外DataFrame.select()一個最妙的使用情境是調整欄位順序，例如：

(df.select("random", "nrs"))
(df.select(["random", "nrs"]))

shape: (5, 2)
┌──────────┬─────┐
│ random   ┆ nrs │
│ ---      ┆ --- │
│ f64      ┆ i64 │
╞══════════╪═════╡
│ 0.37454  ┆ 1   │
│ 0.950714 ┆ 2   │
│ 0.731994 ┆ 3   │
│ 0.598658 ┆ 4   │
│ 0.156019 ┆ 5   │
└──────────┴─────┘

欄位名稱可以一個一個輸入，也可以放在list中一次傳入。

行

DataFrame.filter()

DataFrame.filter()可以讓我們給定條件來選擇行數。

舉例來說，如果只想選取「"nrs"」欄位中大於「"random"」乘10的行，可以這麼寫：

(
    df.filter(pl.col("nrs").gt(pl.col("random").mul(10)))
)

shape: (1, 2)
┌─────┬──────────┐
│ nrs ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 5   ┆ 0.156019 │
└─────┴──────────┘

這裡我們使用了Expr.gt()這個expression來判斷「"nrs"」欄位中的各行之值是否大於後面expression計算出來的各行之值，接著再由DataFrame.filter()這個context來選擇需要回傳的行數。

在DataFrame.filter()中也可以進行邏輯運算，像是只想選取「"nrs"」欄位大於3或是「"nrs"」欄位等於1的行，可以這麼寫：

(
    df.filter(pl.col("nrs").gt(3).or_(pl.col("nrs").eq(1)))
)

shape: (3, 2)
┌─────┬──────────┐
│ nrs ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 1   ┆ 0.37454  │
│ 4   ┆ 0.598658 │
│ 5   ┆ 0.156019 │
└─────┴──────────┘

這裡我們使用了Expr.or_()這個expression來進行bitwise的運算。

分組聚合

DataFrame.group_by().agg()

DataFrame.group_by().agg()可以讓我們針對不同欄位進行不同的聚合操作。

假設我們新增一欄「"group"」欄位到df中，前面三行為「"a"」，而後面兩行為「"b"」。此時，如果想針對「"group"」欄位進行分組，並求出「"nrs"」欄位的總和及「"random"」欄位的平均值，可以這麼寫：

(
    df.with_columns(group=pl.Series(list("aaabb")).explode())
    .group_by("group")
    .agg(pl.col("nrs").sum(), pl.col("random").mean())
)

shape: (2, 3)
┌───────┬─────┬──────────┐
│ group ┆ nrs ┆ random   │
│ ---   ┆ --- ┆ ---      │
│ str   ┆ i64 ┆ f64      │
╞═══════╪═════╪══════════╡
│ b     ┆ 9   ┆ 0.377339 │
│ a     ┆ 6   ┆ 0.685749 │
└───────┴─────┴──────────┘

這裡我們使用了Expr.sum()及Expr.mean()兩種expression。

實務上，使用pl.sum()及pl.mean()這樣的語法糖寫法，也是十分常見的。例如：

(
    df.with_columns(group=pl.Series(list("aaabb")).explode())
    .group_by("group")
    .agg(pl.sum("nrs"), pl.mean("random"))
)

使用[]來存取行或欄

Polars其實也支援以[]的方式來存取行或欄。舉例來說，假如想取得「"random"」及「"nrs"」欄位的第二及第四行，可以使用索引值的方式來取得：

df[[1, 3], [1, 0]] 

shape: (2, 2)
┌──────────┬─────┐
│ random   ┆ nrs │
│ ---      ┆ --- │
│ f64      ┆ i64 │
╞══════════╪═════╡
│ 0.950714 ┆ 2   │
│ 0.598658 ┆ 4   │
└──────────┴─────┘

甚至可以使用類似早期Pandas DataFrame的.ix()（花式索引），在針對行時使用索引值，而在針對欄位時使用欄位名稱，例如：

df[[1, 3], ["random", "nrs"]] 

shape: (2, 2)
┌──────────┬─────┐
│ random   ┆ nrs │
│ ---      ┆ --- │
│ f64      ┆ i64 │
╞══════════╪═════╡
│ 0.950714 ┆ 2   │
│ 0.598658 ┆ 4   │
└──────────┴─────┘

不過我個人覺得這只是一個快速prototyping的捷徑，實務上很少看到有人這樣操作。

備註

註1：建議您也可以參考Michael所寫的部落格文章，其簡明扼要地說明了如何使用Polars搭配gt來快速製作表格。

註2：Polar提供有DataFrame級別的DataFrame.pipe()及expression級別的Expr.pipe()。這兩個功能使得我們容易串接各種不同的context或expression。

註3：實務上，您可以會想要使用更高效的LazyFrame。我們這邊為了方便說明，使用了一般的DataFrame。

Code

本日程式碼傳送門。