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

この記事はPolars Advent Calendar 2023の18日目の記事です！
トミーさん、素晴らしい企画ありがとうございます！このAdvent Calendarで自分もPolarsを学びます👀

はじめに

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

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

ネスト型データ

ネスト型のデータとは

自分が扱ったデータはネスト型のデータが多く、Polarsのネスト型データをよく使いました。
ネスト型データを使用したいデータの例として以下のようなものがあります。

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

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

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は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を扱っていました。

例を見てみましょう

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になっていますね。

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型のデータを作っていきます。

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型のカラムを分解する作業を行うことが多かったです。

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

参考