Polarsでネスト型のデータを扱う

この記事は[Polars Advent Calendar 2023](https://qiita.com/advent-calendar/2023/polars)の18日目の記事です！
[トミーさん](https://qiita.com/Johannyjm)、素晴らしい企画ありがとうございます！このAdvent Calendarで自分もPolarsを学びます👀

# はじめに
私は最近、インターン先で非常に大きなデータをDataFrame形式で操作する必要があり、PandasとPolarsで少し触ったところPolarsの方が高速だったため、良い機会だしPolarsを使おうと思いました。
その際にネスト型のデータを扱うことが多かったので、調べたことを自分なりにまとめます！　

:::note warn
あくまで、個人的なまとめであり正しい保証はないです
:::


# ネスト型データ
### ネスト型のデータとは
自分が扱ったデータは[ネスト型](https://wa3.i-3-i.info/word1284.html)のデータが多く、Polarsのネスト型データをよく使いました。
ネスト型データを使用したいデータの例として以下のようなものがあります。

- 機械学習分野でのembedding
- ネストされたjson
- group_byした結果をリスト形式で保存

以下はgroup_byした結果をリスト形式で保存したいときの例です。

```python
df_pl = pl.DataFrame(
    {
        'id': [1, 1, 2, 2, 2, 3, 3],
        'belonging': ['apple', 'grapefruit', 'apple', 'pencil', 'book', 'pencil', 'book']
    }
)

df_nested = df_pl.group_by('id').agg(
    pl.col('belonging')
    .map_elements(list)
    .alias('belongings')
)

df_nested
"""
(3, 2)
┌─────┬─────────────────────────────┐
│ id  ┆ belongings                  │
│ --- ┆ ---                         │
│ i64 ┆ list[str]                   │
╞═════╪═════════════════════════════╡
│ 1   ┆ ["apple", "card"]           │
│ 2   ┆ ["apple", "pencil", "book"] │
│ 3   ┆ ["pencil", "book"]          │
└─────┴─────────────────────────────┘
"""
```

## Polarsによるネスト型
Polarsでネスト型を扱う[Data types](https://pola-rs.github.io/polars/py-polars/html/reference/datatypes.html)は4つあります
- `pl.List`
- `pl.Array`
- `pl.Object`
- `pl.Struct`

また、この4つを以下の2つのグループに分けることができます
- 各行にシーケンスを格納するグループ: `pl.List`, `pl.Array`, `pl.Object`
- 列のネストされたコレクション: `pl.Struct`

### シーケンスグループ
このシーケンスグループには`pl.List`, `pl.Array`, `pl.Object`が該当しますが、この中でも色々と格納方法などが異なります。
まず、
- `pl.List`, `pl.Array`はPolarsの`Series`として格納
- `pl.Object`はPythonの`list`として格納

という違いがあります。

PolarsのSeriesは同じdtypeを持たなくてならないので、`pl.List`, `pl.Array`となるネスト型のデータは同じdtypeを持っていることになります。
一方、`pl.Object`はPythonのlistとして格納されるので、文字列や数値といった複数のtypeを持つデータをネスト型として格納できます。

ちなみに`pl.List`と `pl.Array`の違いは`pl.List`ではシーケンスの長さが可変長ですが、`pl.Array`固定長です。
~~実際、私はpl.Arrayを使用する機会がなかったため、その特長や利点について詳細に説明することはできません。基本`pl.List`を扱っていました。~~

例を見てみましょう
```python
df_pl  = pl.DataFrame(
    {
        'list': [[0.1, 1], [2, 3]],
        'object_and_num': [['a', 0], ['b', 1]]
    }
).with_columns(array=pl.col('list').cast(pl.Array(width=2, inner=pl.Int64)))

df_pl
"""
(2, 3)
┌────────────┬────────────────┬───────────────┐
│ list       ┆ object_and_num ┆ array         │
│ ---        ┆ ---            ┆ ---           │
│ list[f64]  ┆ object         ┆ array[i64, 2] │
╞════════════╪════════════════╪═══════════════╡
│ [0.1, 1.0] ┆ ['a', 0]       ┆ [0, 1]        │
│ [2.0, 3.0] ┆ ['b', 1]       ┆ [2, 3]        │
└────────────┴────────────────┴───────────────┘
"""
```
上記のコードでは`pl.List`を作成後、`with_columns`や`cast`を利用して`pl.Array`を作成しました。typeを見てみるとseriesとlistになっていますね。

```python
print(type(df_pl['list'][0]))
print(type(df_pl['array'][0]))
print(type(df_pl['object_and_num'][0]))
# <class 'polars.series.series.Series'>
# <class 'polars.series.series.Series'>
# <class 'list'>
```


### コレクション
コレクションには`pl.Struct`が該当します。こちらはネストされたデータが入っているjsonなどを用いる際に使います。

まず、擬似的に`pl.struct`型のデータを作っていきます。
```python
df_struct = pl.DataFrame(
    {
        'num': [1, 2, 3],
        'object': ['a', 'i', 'u'],
        'bool': [True, False, None],
        'list': [[1, 2], [3, 4], [5]]
    }
).select(pl.struct(pl.all()).alias('struct_col'))
"""
shape: (3, 1)
┌──────────────────────┐
│ struct_col           │
│ ---                  │
│ struct[4]            │
╞══════════════════════╡
│ {1,"a",true,[1, 2]}  │
│ {2,"i",false,[3, 4]} │
│ {3,"u",null,[5]}     │
└──────────────────────┘
"""
```
自分はネストされたカラムを含んだjsonを読み込み、struct型のカラムを分解する作業を行うことが多かったです。
```python
df_struct.unnest('struct_col')

"""
(3, 4)
┌─────┬────────┬───────┬───────────┐
│ num ┆ object ┆ bool  ┆ list      │
│ --- ┆ ---    ┆ ---   ┆ ---       │
│ i64 ┆ str    ┆ bool  ┆ list[i64] │
╞═════╪════════╪═══════╪═══════════╡
│ 1   ┆ a      ┆ true  ┆ [1, 2]    │
│ 2   ┆ i      ┆ false ┆ [3, 4]    │
│ 3   ┆ u      ┆ null  ┆ [5]       │
└─────┴────────┴───────┴───────────┘
"""
```

その他にも`struct.field(name)`や`struct.json_encode()`などがあります。
[Polars API reference Struct](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/struct.html)



### 参考

https://pola.rs/

https://pola-rs.github.io/polars/api/

https://www.rhosignal.com/posts/nested-dtypes/#the-plstruct-type-of-nested-columns