Polarsのリスト型に対するよくある操作

Polarsにはリスト型があります。
一見ややこしい処理も、Polarsならではの直感的な操作が使えて便利です。

最近使うことが多かったので、ややニッチかとは思いますが備忘がてら紹介します。

準備

!wget https://jaqket.s3.ap-northeast-1.amazonaws.com/data/aio_01/dev2_questions.json

import polars as pl
pl.Config(fmt_str_lengths=50)
df = pl.read_ndjson('/content/dev2_questions.json')

今回は日本語のクイズのデータを使って紹介します¹²。
クイズの択問題の選択肢がリスト列として格納されています。

基本的なlist操作

要素数

df.select(
    pl.col('answer_candidates'),
    # リストの要素数（選択肢の数）
    pl.col('answer_candidates').list.len().alias('answer_candidates_len')
).head(2)

20択問題とは、、選択肢が多いですね笑

文字列に結合

df.select(
    pl.col('answer_candidates').list.join('/')
).head(2)

縦に展開

df.select('question','answer_candidates').explode('answer_candidates').head(3)

本記事で紹介する操作のなかだと、これだけはデータフレームのshapeが変わる操作です。

あとexplodeの標準的な日本語が分かりません。

要素の抽出

いくつかやり方があるので、それぞれ記載します。

n番目の要素を抽出

df.select(
    pl.col('answer_candidates').list.get(0).alias('get'),
    pl.col('answer_candidates').list.first().alias('first'),
    pl.col('answer_candidates').list[0].alias('slice'),
).head(2)

部分リストを抽出

df.select(
    # 先頭から抽出
    pl.col('answer_candidates').list.head(3).alias('head'),
    # 任意の箇所から抽出
    pl.col('answer_candidates').list.slice(1,3).alias('slice'),
).head(2)

sliceは、start_indexと取得する長さを指定します。

なお、特定の条件に該当する要素だけを抽出する方法に関しては、次節で解説しています。

各要素に対する処理

df.select(
    # ある要素を含むか
    pl.col('answer_candidates').list.contains('タコブネ').alias('タコブネ'),
    pl.col('answer_candidates').list.contains('タコブ').alias('タコブ'),
    # 各要素について、ある文字列（タコブ）を含むか
    pl.col('answer_candidates').list.eval(pl.element().str.contains('タコブ')).alias('タコブ_each'),
    # ある文字列（タコブ）を含む要素が一つでも存在するか
    pl.col('answer_candidates').list.eval(pl.element().str.contains('タコブ')).list.any().alias('タコブ_any'),
).head(2)

eval と pl.element() を組み合わせて使うことで、それぞれの要素に対する処理をほぼなんでも実現できます。

、、、ここまでくるとタコブネがどんな生き物なのか気になってきますよね。こいつらしいです³。

該当する要素だけを抽出

df.select(
    # ひらがなを含む選択肢があるか
    pl.col('answer_candidates').list.eval(
        pl.element().str.contains('[あ-ん]')
    ).list.any().alias('ひらがな'),
    # ひらがなを含む選択肢を抽出
    pl.col('answer_candidates').list.eval(
        pl.element().filter(pl.element().str.contains('[あ-ん]'))
    )
).head(5)

普通のExpressionと似た感じで filter が使えます。

ネストされたリストに対する操作

少し複雑な処理として、リストの要素がリストである場合の操作も同じ考え方で実現できます⁴。

特に意味を持つ例ではないですが、

• 清浄の園 → [清浄、園]
• 宇宙の終焉 → [宇宙、終焉]

みたいに「の」で分割したものを対象に操作してみます。

df.with_columns(
    # 「の」でリストの各要素をリストに分解
    pl.col('answer_candidates').list.eval(
        pl.element().str.split('の')
    )
).select(
    pl.col('question'),
    # 要素数が２以上で、２つ目の要素が一文字の選択肢のみ抽出
    # 「〇〇の✕」に相当する選択肢
    pl.col('answer_candidates').list.eval(
        pl.element().filter(
            (pl.element().list.len() > 1) & 
            (pl.element().list.get(1).str.len_chars() == 1)
        )
    )
).slice(40, 5)

この処理自体は特段listを使わなくても実現できるものですが、このように
pl.element().list …と書くことで、ネストされたリストに対しても同じ考え方で処理を書き続けることが可能になります。

---

ということで、いかがだったでしょうか。

filter や when など、通常のExpressionと似た感覚で操作でき、慣れるハードルは低そうです。Polars最高！！

---

参考

今回はリストの中身が文字列のケースを見ていましたが、 diff や mean など数値に対する処理も豊富にあります。そのあたり詳細はドキュメントをぜひ。

1. JAQKET: クイズを題材にした日本語QAデータセット。https://www.nlp.ecei.tohoku.ac.jp/projects/jaqket/ ↩

2. 鈴木正敏, 鈴木潤, 松田耕史, ⻄田京介, 井之上直也. “JAQKET:クイズを題材にした日本語QAデータセットの構築”. 言語処理学会第26回年次大会(NLP2020) 発表論文集 ↩

3. 画像はWikipediaより引用。https://ja.wikipedia.org/wiki/%E3%82%BF%E3%82%B3%E3%83%96%E3%83%8D ↩

4. 複雑なので、このやりかたが良いかは微妙なところです。さっさとexplodeして処理するといった方法もきっとあります。 ↩