想要研究命名空間的動機來自於，在[Day23]及[Day24]使用Plotnine及Great Tables製作關稅表時，會重覆用到如何將奇數及偶數列填上不同顏色的邏輯。原先只是想將該邏輯抽取為binarize()，但後來發現可以使用命名空間，將相關邏輯封裝在其中。

Polars針對pl.Series、pl.LazyFrame、pl.DataFrame及pl.Expr提供了客製化命名空間的功能。我們會以pl.Expr為例，一起客製化spt命名空間。其將提供三個函數，幫助使用者依照各行的索引值進行分組：

• binarize()：接收兩個變數，分別做為經過mod 2運算後的填充值。
• trinarize()：接收三個變數，分別做為經過mod 3運算後的填充值。
• bucketize()：接收一個列表，其內含n個變數，分別做為經過mod n運算後的填充值。

本日大綱如下：

0. 本日引入模組及準備工作
1. 註冊命名空間
2. 輔助函數_mod_expr()
3. binarize()：分為兩組
4. trinarize()：分為三組
5. bucketize()：分為n組
6. 實例應用

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

from typing import Any

import polars as pl
from great_tables import GT, loc, style

1. 註冊命名空間

建立DiscreteSplitter class，並使用pl.api.register_expr_namespace()做為裝飾器。裝飾器內的參數即為此命名空間之名，此處定義為「"spt"」。

@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def __init__(self, expr: pl.Expr) -> None:
        self._expr = expr

2. 輔助函數_mod_expr()

如果想產生一列連續數字做為索引列，一般會使用pl.DataFrame.with_row_index()，但其為一個pl.DataFrame級別的函數。所幸在該文件的最下方，提供了pl.Expr級別的公式。我們將此公式後接上pl.Expr.mod()運算，封裝於_mod_expr()內，方便重覆呼叫。

@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def _mod_expr(self, n: int) -> pl.Expr:
        return pl.int_range(pl.len(), dtype=pl.UInt32).mod(n)

3.binarize()：分為兩組

binarize()利用pl.when().then().otherwise()進行分組，將新增列之奇數行值設定為lit1，偶數行則設定為lit2。

@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def binarize(
        self, lit1: str, lit2: str, name: str = "binarized"
    ) -> pl.Expr:
        mod_expr = self._mod_expr(2)
        return (
            pl.when(mod_expr.eq(0))
            .then(pl.lit(lit1))
            .otherwise(pl.lit(lit2))
            .alias(name)
        )

4. trinarize()：分為三組

trinarize()的架構與binarize()一樣，只是此處分為三組，所以多增加了一組pl.when().then()的判斷式。

@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def trinarize(
        self, lit1: str, lit2: str, lit3: str, name: str = "trinarized"
    ) -> pl.Expr:
        mod_expr = self._mod_expr(3)
        return (
            pl.when(mod_expr.eq(0))
            .then(pl.lit(lit1))
            .when(mod_expr.eq(1))
            .then(pl.lit(lit2))
            .otherwise(pl.lit(lit3))
            .alias(name)
        )

5. bucketize()：分為n組

bucketize是一個通用的分組寫法，此處我們改為接收列表做為參數，其中的元素個數為分組組數，而其內元素則為分組後之填充值。

@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def bucketize(
        self, lits: list[Any], name: str = "bucketized"
    ) -> pl.Expr:
        mod_expr = self._mod_expr(len(lits))

        # first
        expr = pl.when(mod_expr.eq(0)).then(pl.lit(lits[0]))

        # middles
        for i, one_lit in enumerate(lits[1:-1], start=1):
            expr = expr.when(mod_expr.eq(i)).then(pl.lit(one_lit))

        # last
        expr = expr.otherwise(pl.lit(lits[-1]))
        return expr.alias(name)

6. 實例應用

以下展示如何使用spt命名空間中的三個函數來新增三列：

df = (
    pl.DataFrame({"n": [100, 50, 72, 83, 97, 42, 20, 51, 77]})
    .with_row_index(offset=1)
    .with_columns(
        pl.col("").spt.binarize("lightblue", "papayawhip"),
        pl.col("").spt.trinarize("one", "two", "three"),
        pl.col("").spt.bucketize([1, 2, 3, 4]),
    )
)

shape: (9, 5)
┌───────┬─────┬────────────┬────────────┬────────────┐
│ index ┆ n   ┆ binarized  ┆ trinarized ┆ bucketized │
│ ---   ┆ --- ┆ ---        ┆ ---        ┆ ---        │
│ u32   ┆ i64 ┆ str        ┆ str        ┆ i32        │
╞═══════╪═════╪════════════╪════════════╪════════════╡
│ 1     ┆ 100 ┆ lightblue  ┆ one        ┆ 1          │
│ 2     ┆ 50  ┆ papayawhip ┆ two        ┆ 2          │
│ 3     ┆ 72  ┆ lightblue  ┆ three      ┆ 3          │
│ 4     ┆ 83  ┆ papayawhip ┆ one        ┆ 4          │
│ 5     ┆ 97  ┆ lightblue  ┆ two        ┆ 1          │
│ 6     ┆ 42  ┆ papayawhip ┆ three      ┆ 2          │
│ 7     ┆ 20  ┆ lightblue  ┆ one        ┆ 3          │
│ 8     ┆ 51  ┆ papayawhip ┆ two        ┆ 4          │
│ 9     ┆ 77  ┆ lightblue  ┆ three      ┆ 1          │
└───────┴─────┴────────────┴────────────┴────────────┘

請留意，由於spt命名空間中沒有使用到self._expr，所以可以使用任意的expr，例如此例中的pl.col("")。

參考資料

個人部落格文章：Polars Custom Expression Namespace。

Code

本日程式碼傳送門。