Polars によるデータ加工テクニック

思いついたら追記しよう

## 列をまとめて選択

公式ドキュメント
* [More flexible column selections / Expression expansion - Polars user guide](https://docs.pola.rs/user-guide/expressions/expression-expansion/#more-flexible-column-selections)
* [Selectors — Polars documentation](https://docs.pola.rs/api/python/stable/reference/selectors.html)

例
* 文字列の列を NFKC で Unicode 正規化して前後の空白を削除
  `df.with_columns(cs.string().str.normalize("NFKC").str.strip_chars())`
* 途中計算の列名をアンダースコア始まりとしておいて出力前にまとめて削除
  `df.drop(cs.starts_with("_"))`
* altair でグラフを書こうとしたら duration が対応していなかったのでまとめて削除
  `df.drop(cs.duration())`
* やっぱり分単位の経過時間として残したくなったのでまとめて計算
  `df.with_columns(cs.duration().dt.total_seconds() / 60)`

## 列名を一括変更

公式ドキュメント
* [polars.Expr.name.map — Polars documentation](https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.name.map.html)

例
* スプレッドシートへ出力するときに `_` 区切りの列名を改行区切りにする
  `df.with_columns(pl.all().name.map(lambda x: x.replace("_", "\n")))`

## 横持ちタプルを縦持ちタプルに整形

Excel や DB でやりがちな `(key, value)` の固定長配列を横持ちとする構造を Polars で集計しやすい縦持ちに変形する。

公式ドキュメント
* [polars.concat_list — Polars documentation](https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.concat_list.html)
* [polars.DataFrame.explode — Polars documentation](https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.explode.html)

### 整形前
```plaintext
shape: (10, 7)
┌───────┬───────┬─────────┬───────┬─────────┬───────┬─────────┐
│ index ┆ key01 ┆ value01 ┆ key02 ┆ value02 ┆ key03 ┆ value03 │
│ ---   ┆ ---   ┆ ---     ┆ ---   ┆ ---     ┆ ---   ┆ ---     │
│ i64   ┆ i64   ┆ f64     ┆ i64   ┆ f64     ┆ i64   ┆ f64     │
╞═══════╪═══════╪═════════╪═══════╪═════════╪═══════╪═════════╡
│ 0     ┆ 0     ┆ 0.0     ┆ 10    ┆ 0.0     ┆ 20    ┆ 0.0     │
│ 1     ┆ 1     ┆ 1.0     ┆ 11    ┆ 2.0     ┆ 21    ┆ 3.0     │
│ 2     ┆ 2     ┆ 2.0     ┆ 12    ┆ 4.0     ┆ 22    ┆ 6.0     │
│ 3     ┆ 3     ┆ 3.0     ┆ 13    ┆ 6.0     ┆ 23    ┆ 9.0     │
│ 4     ┆ 4     ┆ 4.0     ┆ 14    ┆ 8.0     ┆ 24    ┆ 12.0    │
│ 5     ┆ 5     ┆ 5.0     ┆ 15    ┆ 10.0    ┆ 25    ┆ 15.0    │
│ 6     ┆ 6     ┆ 6.0     ┆ 16    ┆ 12.0    ┆ 26    ┆ 18.0    │
│ 7     ┆ 7     ┆ 7.0     ┆ 17    ┆ 14.0    ┆ 27    ┆ 21.0    │
│ 8     ┆ 8     ┆ 8.0     ┆ 18    ┆ 16.0    ┆ 28    ┆ 24.0    │
│ 9     ┆ 9     ┆ 9.0     ┆ 19    ┆ 18.0    ┆ 29    ┆ 27.0    │
└───────┴───────┴─────────┴───────┴─────────┴───────┴─────────┘
```

### 整形中 (`select` まで)
```plaintext
shape: (10, 4)
┌───────┬────────────────────┬─────────────┬───────────────────┐
│ index ┆ _id                ┆ key         ┆ value             │
│ ---   ┆ ---                ┆ ---         ┆ ---               │
│ i64   ┆ list[str]          ┆ list[i64]   ┆ list[f64]         │
╞═══════╪════════════════════╪═════════════╪═══════════════════╡
│ 0     ┆ ["01", "02", "03"] ┆ [0, 10, 20] ┆ [0.0, 0.0, 0.0]   │
│ 1     ┆ ["01", "02", "03"] ┆ [1, 11, 21] ┆ [1.0, 2.0, 3.0]   │
│ 2     ┆ ["01", "02", "03"] ┆ [2, 12, 22] ┆ [2.0, 4.0, 6.0]   │
│ 3     ┆ ["01", "02", "03"] ┆ [3, 13, 23] ┆ [3.0, 6.0, 9.0]   │
│ 4     ┆ ["01", "02", "03"] ┆ [4, 14, 24] ┆ [4.0, 8.0, 12.0]  │
│ 5     ┆ ["01", "02", "03"] ┆ [5, 15, 25] ┆ [5.0, 10.0, 15.0] │
│ 6     ┆ ["01", "02", "03"] ┆ [6, 16, 26] ┆ [6.0, 12.0, 18.0] │
│ 7     ┆ ["01", "02", "03"] ┆ [7, 17, 27] ┆ [7.0, 14.0, 21.0] │
│ 8     ┆ ["01", "02", "03"] ┆ [8, 18, 28] ┆ [8.0, 16.0, 24.0] │
│ 9     ┆ ["01", "02", "03"] ┆ [9, 19, 29] ┆ [9.0, 18.0, 27.0] │
└───────┴────────────────────┴─────────────┴───────────────────┘
```

### 整形後
```plaintext
shape: (30, 4)
┌───────┬─────┬─────┬───────┐
│ index ┆ _id ┆ key ┆ value │
│ ---   ┆ --- ┆ --- ┆ ---   │
│ i64   ┆ str ┆ i64 ┆ f64   │
╞═══════╪═════╪═════╪═══════╡
│ 0     ┆ 01  ┆ 0   ┆ 0.0   │
│ 0     ┆ 02  ┆ 10  ┆ 0.0   │
│ 0     ┆ 03  ┆ 20  ┆ 0.0   │
│ 1     ┆ 01  ┆ 1   ┆ 1.0   │
│ 1     ┆ 02  ┆ 11  ┆ 2.0   │
│ …     ┆ …   ┆ …   ┆ …     │
│ 8     ┆ 02  ┆ 18  ┆ 16.0  │
│ 8     ┆ 03  ┆ 28  ┆ 24.0  │
│ 9     ┆ 01  ┆ 9   ┆ 9.0   │
│ 9     ┆ 02  ┆ 19  ┆ 18.0  │
│ 9     ┆ 03  ┆ 29  ┆ 27.0  │
└───────┴─────┴─────┴───────┘
```

### コード
```py
import numpy as np
import polars as pl
import polars.selectors as cs

df = pl.DataFrame({
    "index": np.arange(10),
    "key01": np.arange(10),
    "value01": 1.0 * np.arange(10),
    "key02": np.arange(10, 20),
    "value02": 2.0 * np.arange(10),
    "key03": np.arange(20, 30),
    "value03": 3.0 * np.arange(10),
})

(
    df
    .select(
        "index",
        # NOTE: 列がソート済みと仮定している
        _id=pl.lit(df.select(cs.starts_with("key")).columns).list.eval(
            pl.element().str.extract(r"(\d+)")
        ),
        key=pl.concat_list(cs.starts_with("key")),
        value=pl.concat_list(cs.starts_with("value")),
    )
    .explode("_id", "key", "value")
)
```

<details><summary>key と value で型が同じなら使えるやつ</summary>

`unpivot` の時点で `key` と `value` が同じ列になってしまうため、型昇格ができない組み合わせだと死ぬ

```py
import numpy as np
import polars as pl
import polars.selectors as cs

df = pl.DataFrame({
    "index": np.arange(10),
    "key01": np.arange(10),
    "value01": 1 * np.arange(10),
    "key02": np.arange(10, 20),
    "value02": 2 * np.arange(10),
    "key03": np.arange(20, 30),
    "value03": 3 * np.arange(10),
})

(
    df.unpivot(
        cs.starts_with("key", "value"),
        index=["index"],
        variable_name="__variable",
        value_name="__value",
    )
    .with_columns(
        pl.col("__variable")
        .str.extract_groups(r"^(?<_name>.*?)(?<_id>\d+)$")
        .struct.unnest()
    )
    .pivot("_name", index=["index", "_id"], values="__value")
    .sort("index", "_id")
)
```
</details>

## 自己結合

[`polars.DataFrame.pipe`](https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.pipe.html) を使うと再代入や一次変数をなくせる。

### 整形前
```plaintext
shape: (9, 3)
┌───────┬────────┬───────┐
│ index ┆ index2 ┆ value │
│ ---   ┆ ---    ┆ ---   │
│ i64   ┆ i64    ┆ f64   │
╞═══════╪════════╪═══════╡
│ 1     ┆ 1      ┆ 0.0   │
│ 1     ┆ 2      ┆ 2.0   │
│ 1     ┆ 3      ┆ 4.0   │
│ 2     ┆ 1      ┆ 6.0   │
│ 2     ┆ 2      ┆ 8.0   │
│ 2     ┆ 3      ┆ 10.0  │
│ 3     ┆ 1      ┆ 12.0  │
│ 3     ┆ 2      ┆ 14.0  │
│ 3     ┆ 3      ┆ 16.0  │
└───────┴────────┴───────┘
```

### 整形後
```plaintext
shape: (9, 4)
┌───────┬────────┬───────┬────────────┐
│ index ┆ index2 ┆ value ┆ value_diff │
│ ---   ┆ ---    ┆ ---   ┆ ---        │
│ i64   ┆ i64    ┆ f64   ┆ f64        │
╞═══════╪════════╪═══════╪════════════╡
│ 1     ┆ 1      ┆ 0.0   ┆ null       │
│ 1     ┆ 2      ┆ 2.0   ┆ 2.0        │
│ 1     ┆ 3      ┆ 4.0   ┆ 2.0        │
│ 2     ┆ 1      ┆ 6.0   ┆ null       │
│ 2     ┆ 2      ┆ 8.0   ┆ 2.0        │
│ 2     ┆ 3      ┆ 10.0  ┆ 2.0        │
│ 3     ┆ 1      ┆ 12.0  ┆ null       │
│ 3     ┆ 2      ┆ 14.0  ┆ 2.0        │
│ 3     ┆ 3      ┆ 16.0  ┆ 2.0        │
└───────┴────────┴───────┴────────────┘
```

### コード

#### 複雑な計算を入れられるやつ
```py
df = pl.DataFrame({
    "index": [1] * 3 + [2] * 3 + [3] * 3,
    "index2": [1, 2, 3] * 3,
    "value": 2.0 * np.arange(9),
})

(
    df.pipe(
        lambda _df: _df.join(
            _df,
            how="left",
            left_on=["index", "index2"],
            right_on=["index", pl.col("index2") + 1],
        )
    )
    .with_columns(value_diff=pl.col("value") - pl.col("value_right"))
    .drop(cs.ends_with("_right"))
)
```

#### 複雑な計算をしない場合
```py
df = pl.DataFrame({
    "index": [1] * 3 + [2] * 3 + [3] * 3,
    "index2": [1, 2, 3] * 3,
    "value": 2.0 * np.arange(9),
})

df.with_columns(
    value_diff=pl.col("value").diff().over("index", order_by="index2"),
)
```

## 腐った ON, OFF 信号の前処理

### コード
```py
import matplotlib.pyplot as plt
import polars as pl

df = pl.DataFrame({
    #          x, x, x, x        x, x, x, x
    "on":  [0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0],
    "off": [1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1],
}).with_row_index()

df_out = (
    df.with_columns(
        # ON, OFF 信号を立ち上がりパルスだけ残す
        pl.col("on", "off").cast(pl.Int64).diff().clip(0, 1).cast(pl.Boolean)
        .name.suffix("_pulse")
    )
    .with_columns(
        # ON の立ち上がりの時は ON (OFF 立ち上がりでも無視)
        state=pl.when("on_pulse").then(True)
        # OFF 立ち上がり単独の時は OFF
        .when("off_pulse").then(False)
        # それ以外は前の状態から継続
        .otherwise(None)
        .fill_null(strategy="forward")
    )
    # 初期状態だけ定義できないので OFF を与える
    .fill_null(False)
)

fig, ax = plt.subplots(nrows=5, layout="constrained")
ax[0].step("index", "on", where="post", data=df_out)
ax[0].set_ylabel("ON")
ax[1].step("index", "off", where="post", data=df_out)
ax[1].set_ylabel("OFF")
ax[2].step("index", "on_pulse", where="post", data=df_out)
ax[2].set_ylabel("ON Pulse")
ax[3].step("index", "off_pulse", where="post", data=df_out)
ax[3].set_ylabel("OFF Pulse")
ax[4].step("index", "state", where="post", data=df_out)
ax[4].set_ylabel("ON State")
fig.show()
```

### 結果
![Figure_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64912/fdc50aa4-2a19-4df8-941a-387b39197849.png)

## 指定値が出てくるまで選択

特定のイベントが来たら後のデータはいらないときに

```py
import numpy as np
import polars as pl

df = pl.DataFrame({
    "index": np.arange(8),
    "event": [0, 2, 100, 0, 42, 0, 100, 3],
    "value": np.linspace(0.1, 0.8, 8),
})

df.filter((pl.col("event") == 42).rle_id() <= 1)
```

```plaintext
shape: (8, 3)
┌───────┬───────┬───────┐
│ index ┆ event ┆ value │
│ ---   ┆ ---   ┆ ---   │
│ i64   ┆ i64   ┆ f64   │
╞═══════╪═══════╪═══════╡
│ 0     ┆ 0     ┆ 0.1   │
│ 1     ┆ 2     ┆ 0.2   │
│ 2     ┆ 100   ┆ 0.3   │
│ 3     ┆ 0     ┆ 0.4   │
│ 4     ┆ 42    ┆ 0.5   │
│ 5     ┆ 0     ┆ 0.6   │
│ 6     ┆ 100   ┆ 0.7   │
│ 7     ┆ 3     ┆ 0.8   │
└───────┴───────┴───────┘
```

```plaintext
shape: (5, 3)
┌───────┬───────┬───────┐
│ index ┆ event ┆ value │
│ ---   ┆ ---   ┆ ---   │
│ i64   ┆ i64   ┆ f64   │
╞═══════╪═══════╪═══════╡
│ 0     ┆ 0     ┆ 0.1   │
│ 1     ┆ 2     ┆ 0.2   │
│ 2     ┆ 100   ┆ 0.3   │
│ 3     ┆ 0     ┆ 0.4   │
│ 4     ┆ 42    ┆ 0.5   │
└───────┴───────┴───────┘
```