この記事は朝日新聞社Advent Calendar2024の11日目の記事です。
昨日の記事は村瀬さんのAWS Lambda SnapStartを試してみたでした。

Polars 入門向けよく使いそうな機能の逆引きリファレンス

こんにちは、朝日新聞社の新妻です。

皆さん、Polars使ってますか？
自分はこの半年間くらいPandasからPolarsに乗り換えて、しばらく使ってみていました。
個人的な感想として、メモリの効率さや処理の高速さが非常に良くて、特にSNSの投稿のような大規模なデータを扱うときには非常に助かっています。

ということで、非常にオススメできるのですが、
慣れてるツールを乗り換えたりするのって結構ハードルが高いですよね。
ということで、個人的な備忘録も兼ねてPolarsの機能の簡易的な逆引きリファレンスを作ってみました。

ここ半年間の自身の利用履歴から候補を絞っているので、紹介している関数は一部です。
また、使ったPolarsのバージョンは1.2.0から1.10.0あたりです。
（リリース頻度が非常に早くて、今は1.15.0くらい。）

そして、紹介する関数の参照場所は基本的にこれ。
Python API reference — Polars documentation

import

import polars as pl # 慣例的にplにする（Pandasに倣っているのだと思う）

データの読み書き編

詳しくはこちら。
Input/output — Polars documentation
Functions — Polars documentation

以下はよく使いそうなデータなどを抜粋。

CSV / TSV

### 読み出し
df = pl.read_csv(PATH)
df = pl.read_csv(PATH, separator="\t") # TSVの場合
df = pl.read_csv(PATH, has_header=False) # ヘッダーなしの場合
df = pl.read_csv(PATH, encoding="cp932") # UTF-8以外のテキストの場合
df = pl.read_csv(PATH, skip_rows=N) # N行スキップ
df = pl.read_csv(PATH, ignore_errors=True) # parseで発生するエラーを無視する気持ち
df = pl.read_csv(PATH, columns=["col1", "col2"]) # 特定のカラムだけ使いたい時
df = pl.read_csv(PATH, new_columns=["new_name_1", "new_name_2"]) # カラム名の変更

### 書き出し
df.write_csv(PATH)

polars.read_csv — Polars documentation
polars.DataFrame.write_csv — Polars documentation

JSONL / NDJSON

### 読み出し
df = pl.read_ndjson(PATH)
df = pl.read_ndjson(PATH, schema=["new_name_1", "new_name_2"]) # カラム名の変更
df = pl.read_ndjson(PATH, schema={"new_name_1": pl.String, "new_name_2": pl.Int8}) # カラム名の変更 & 型の指定
df = pl.read_ndjson(PATH, ignore_errors=True) # parseで発生するエラーを無視する気持ち

### 書き出し
df = pl.write_ndjson(PATH)

read_ndjsonはネストしているカラムでも読み込めます。
PolarsではStructという型が用意されていて、ネストしているカラムにはそれが当てはめられるます。
このStructの扱いについては、Structを扱うで。

# 例示のためにjsonlをサクッと用意
import io
jsonl_str = '''{"col1": "col1_text", "col2": {"nested_col": "nested_col_text1", "nested_col_list": [1, 2, 3]}}
{"col1": "col1_text", "col2": {"nested_col": "nested_col_text2", "nested_col_list": [4, 5, 6, 7]}}
{"col1": "col1_text", "col2": {"nested_col": "nested_col_text3", "nested_col_list": [8]}}'''
file = io.StringIO(jsonl_str)

pl.read_ndjson(file)
# ┌───────────┬─────────────────────────────────┐
# │ col1      ┆ col2                            │
# │ ---       ┆ ---                             │
# │ str       ┆ struct[2]                       │
# ╞═══════════╪═════════════════════════════════╡
# │ col1_text ┆ {"nested_col_text1",[1, 2, 3]}  │
# │ col1_text ┆ {"nested_col_text2",[4, 5, … 7… │
# │ col1_text ┆ {"nested_col_text3",[8]}        │
# └───────────┴─────────────────────────────────┘

polars.read_ndjson — Polars documentation
polars.DataFrame.write_ndjson — Polars documentation

Excel

# 事前にfastexcelをinstallする必要がある (xlsx2csv, openpyxlでも可だが設定が必要)

### 読み出し
df = pl.read_excel(PATH)
df = pl.read_excel(PATH, sheet_id=1) # シート番号を選ぶ
df = pl.read_excel(PATH, has_header=False) # 一行目がヘッダーじゃない場合
df = pl.read_excel(PATH, columns=["col1", "col2"]) # 読み出すカラム名を指定

### 書き出し
df.write_excel(PATH)

polars.read_excel — Polars documentation
polars.DataFrame.write_excel — Polars documentation

Parquet

### 読み出し
df = pl.read_parquet(PATH)
df = pl.read_parquet(PATH, columns=["col1", "col2"]) # 読み出すカラム名を指定

### 書き出し
df.write_parquet(PATH) # デフォルトはzstd
df.write_parquet(PATH, compression='lz4') # いくつかのcompression方法を選ぶことができる
df.write_parquet(PATH, use_pyarrow=True) # より多くの機能がサポートされているC++実装が使えるらしい

polars.read_parquet — Polars documentation
polars.DataFrame.write_parquet — Polars documentation

Dict / Pandas / datasets など

# 読み出し
df = pl.from_dict({"col_a": [1, 2, 3]}) # Dictから
df = pl.from_records([(1, "a"), (2, "b")], scheme=["col_a", "col_b"]) # Tupleのリストから
df = pl.from_pandas(df_pandas) # Pandasから
df = dataset.to_polars() # HuggingFace datasetsから

# 書き出し
df.rows() # Tupleのリストとして出力
df.rows(named=True) # Dictのリストとして出力
df.to_pandas() # Pandasに変換
dataset = datasets.Dataset(df.to_arrows()) # HuggingFace datasetsに変換

datasetsのデータセットのURLを指定して、Polarsで直接を読むこともできるらしい。（以下を参照）
Polars | Hugging Face

DuckDB

import duckdb
import polars as pl

# PolarsからDuckDB
duckdb.sql("SELECT * FROM df").show() # dfはデータフレームの変数名

# DuckDBからPolars
df = duckdb.sql("SELECT * FROM table" ).pl() # 公式の方法だがレコードが多いとメモリが足りなくなることがある

# メモリが足りないときは、こうするとうまくいくことがある
rows = conn.sql('SELECT * FROM table;').fetchall()
df = pl.DataFrame(rows, schema=["col_a", "col_b"], orient="row")

Integration with Polars – DuckDB

基本操作

カラムの変更・追加をする

# 選択したカラムだけで新しくデータフレームを作りたい。
df.select(
    pl.col("col_a"),
    pl.col("col_b").alias("new_col_b"), # 一応、renameもできる
    # カラムに対する処理の結果をis_thinkingという名前のカラムとして追加
    pl.col("col_c").str.contains("🤔").alias("is_thinking"),
    # aliasなどをつけないと元の名前のカラムを上書きする
    pl.col("col_c").str.contains("🤔"),
)

# 現在のカラムを残したまま、特定のカラムを変更したり、新しいカラムを追加したい
df.with_columns(
    pl.col("col_a"),
    pl.col("col_b").alias("new_col_b"), # 一応、renameもできる
    # カラムに対する処理の結果をis_thinkingという名前のカラムとして追加
    pl.col("col_c").str.contains("🤔").alias("is_thinking"),
    # aliasなどをつけないと元の名前のカラムを上書きする
    pl.col("col_c").str.contains("🤔"),
)

カラムに対する関数の全容はExpressions — Polars documentation

カラムごとに集計(group by)

# col_aの要素でgroup by
df.group_by("col_a").agg(
    # colをそのままagg関数の中に置くと、要素がpl.Listとして保持される
    pl.col("col_b"),
    # もちろんだが集計関数等は用意されている
    pl.col("col_c").n_unique(),
)

# polarsのgroup_byは並列で計算するため、デフォルトでは順序を維持しない
# maintain_orderをTrueにすると順序を維持するが遅くなる
df.group_by("col_a", maintain_order=True).agg(
  pl.col("col_b"),
  pl.col("col_c").n_unique(),
)

詳しくはGroupBy — Polars documentation
経験上だが、group_byの結果をiterateすると基本的に遅い。

ちなみにあるカラムaごとにカラムcの平均を計算して元のデータフレームに追加したい、とかの程度の場合はoverが使える。
overについては特定のgroupごとに要素に対して何らかのを計算したい時で。

条件に合致した行だけを抽出

# filterはTrueになった行だけを抽出する
# col_aに🤔が含まれている行だけを抽出
df.filter(pl.col("col_a").str.contains("🤔"))

polars.Expr.filter — Polars documentation

欠損値を埋める / 含まれている行を落とす

# null/nanのある行を落とす
df.drop_nulls()
df.drop_nulls(subset=["col_a"]) # 特定のカラムのnullを対象
df.drop_nans()
df.drop_nans(subset=["col_a"]) # 特定のカラムのnanを対象

df.fill_nan(0) # nanを0で埋める
df.fill_null(99) # nullを99で埋める
df.fill_null(strategy="zero") # nullを0で埋める
df.fill_null(strategy="forward") # nullを一つ前の値で埋める
df.fill_null(strategy="backward") # nullを一つ後の値で埋める
df.fill_null(strategy="max") # nullを最大値で埋める

fill_nullは他にも埋め方があったりします。
polars.DataFrame.fill_null — Polars documentation

その他

df.describe() # データフレームの要約統計量を一覧にする
df.value_counts("col_a") # col_aの値のカウントを集計
df.value_counts("col_a", normalize=True) # カウントを割合で出せる
df.rename({"col_a": "new_col_a"}) # col_aをnew_col_aにリネーム

カラムの操作

詳しくはこちら。
ここではよく使いそうなものを抜粋。
Expressions — Polars documentation

数値演算など

基本的な操作はほぼできる
- Computation — Polars documentation

# いくつかの例
df.with_columns(
    pl.col("col_a").abs(), # 絶対値を計算
    pl.col("col_a").cum_sum(), # 累積和を計算
    pl.col("col_a").log(), # 　対数変換
    pl.col("col_a").exp(), # 指数関数
    pl.col("col_a").hash(), # ハッシュ化
)

文字列を扱う

Pandasでできるものはほとんどできる。
- String — Polars documentation

# いくつかの例
df.with_columns(
    # 正規表現でマッチしたN番目のパターンを抽出する。
    pl.col("col_a").str.extract(REGEX_PATTERN, group_index=N).alias("pattern_1"),
    # 時間の文字列をPolarsのDateTime型に変換する
    pl.col("col_time_str").str.to_datetime(TIME_FORMAT),
    # strip系はPandasやPythonの標準の名前と違う命名なので注意
    pl.col("col_b").str.strip_chars(" "),
    pl.col("col_b").str.strip_chars_start("- "),
    pl.col("col_b").str.strip_chars_end("。"),
)

JSON文字列が入っているカラム

余談ですが、カラム内のJSON文字列を扱うこともできます。

df.with_columns(
  pl.col("col_json_str").str.json_decode(), # pl.Structに変換、Structについては後述
  pl.col("col_json_str").str.json_path_match("$.a"), # JSON文字列のkey`a`にアクセスする
)

時間を扱う

分析でよく使う基本的なものは揃っています。

全操作はこちら
- Temporal — Polars documentation

df.with_columns(
    # （文字列のところにも書いたが）時間の文字列をPolarsのDateTime型に変換する
    pl.col("col_datetime_str").str.to_datetime(TIME_FORMAT),
    # 逆に時間型を特定の日付フォーマットの文字列に変換する
    pl.col("col_datetime").dt.to_string(TIME_FORMAT),
    # 営業日分の日数追加する
    # 通常は土日だけの考慮だが引数で休日や祝日などをコントロールできる
    pl.col("col_datetime").dt.add_business_days(N),
    # それぞれのカラムに指定した時間を追加する（この例は1年）
    pl.col("col_datetime").dt.offset_by(by: '1y'),
    pl.col("col_datetime").dt.epoch(), # UNIXタイムスタンプを返す
    # もちろんPandasにあったような年を返すようなものは一通り揃っている。
    pl.col("col_datetime").dt.year(),
    pl.col("col_datetime").dt.month(),
    pl.col("col_datetime").dt.week(),
    pl.col("col_datetime").dt.day(),
)

リストを扱う

全操作はこちら
- List — Polars documentation

# よく使いそうな例

# サンプルデータフレーム
df = pl.DataFrame({
    "user": [1, 2, 3],
    "messages": [["text1", "text2", "text3"], ["text4"], ["text5"]]     
})
# ┌──────┬─────────────────────────────┐
# │ user ┆ messages                    │
# │ ---  ┆ ---                         │
# │ i64  ┆ list[str]                   │
# ╞══════╪═════════════════════════════╡
# │ 1    ┆ ["text1", "text2", "text3"] │
# │ 2    ┆ ["text4"]                   │
# │ 3    ┆ ["text5"]                   │
# └──────┴─────────────────────────────┘

# リストの要素を行カラムに展開する
# もちろんだが、対応する他のカラムの要素ごとにコピーされている
df.explode("messages")
# ┌──────┬──────────┐
# │ user ┆ messages │
# │ ---  ┆ ---      │
# │ i64  ┆ str      │
# ╞══════╪══════════╡
# │ 1    ┆ text1    │
# │ 1    ┆ text2    │
# │ 1    ┆ text3    │
# │ 2    ┆ text4    │
# │ 3    ┆ null     │
# └──────┴──────────┘

df.with_columns(
    pl.col("messages").list.gather(0), # リストの0番目の要素だけを取り出す
    # OFFSETからOFFSET+1までのindexの要素を取り出す
    # gatherがindexがないとOutOfBoudsが出るのに対してこちらは出ないため
    # 単体の値を取り出すときにも使い分けで使うことも
    pl.col("messages").list.slice(OFFSET, N),
    # 与えられたリストとの集合演算もできる
    pl.col("messages").list.set_intersection(["text1", "text2"]), # 共通集合のみを取り出す
    pl.col("messages").list.set_union(["text1", "text2"]), # ["text1", "text2"]を足す
    pl.col("messages").list.set_difference(["text1", "text2"]), # ["text1", "text2"]を引く
    # Struct型に変換する
    # 後述するがunnestでリストの個々の要素をカラムに変換したいときに有用
    pl.col("messages").list.to_struct(),
)

Structを扱う

これはPolarsで用意されている辞書型に近い雰囲気。
JSON Linesを読み込んだときに、ネストしているカラムはこの型になっていたりします。

操作はあまり多くない
- Struct — Polars documentation

df = pl.DataFrame({
    "post": [
      {"title": "title1", "post": "post1", "attrs": [3, 1, 2], "user": {"id": 1, "name": "name1"}},
      {"title": "title4", "post": "post4", "attrs": [2, 5, 3], "user":  {"id": 2, "name": "name2"}},
      {"title": "title5", "post": "post5", "attrs": [8, 18, 3], "user":  {"id": 3, "name": "name3"}},
    ],
})

# ┌─────────────────────────────────┐
# │ post                            │
# │ ---                             │
# │ struct[4]                       │
# ╞═════════════════════════════════╡
# │ {"title1","post1",[3, 1, 2],{1… │
# │ {"title4","post4",[2, 5, 3],{2… │
# │ {"title5","post5",[8, 18, 3],{… │
# └─────────────────────────────────┘

df.with_columns(
    pl.col("post").struct.field("title"), # structからtitle要素を取り出す
    pl.col("post").struct.field("user").struct.field("name"), # ちょっと面倒だが深いところも取り出せる
    pl.col("post").struct.json_encode(), # JSON文字列に変換
    pl.col("post").rename_fields(["tweet"]), # fieldの名前(キー)を変更
)

# Structを解体して、キーをカラム名にして、valueを各要素とした行に作り直す
df = df.unnest("post")
# ┌────────┬───────┬────────────┬─────────────┐
# │ title  ┆ post  ┆ attrs      ┆ user        │
# │ ---    ┆ ---   ┆ ---        ┆ ---         │
# │ str    ┆ str   ┆ list[i64]  ┆ struct[2]   │
# ╞════════╪═══════╪════════════╪═════════════╡
# │ title1 ┆ post1 ┆ [3, 1, 2]  ┆ {1,"name1"} │
# │ title4 ┆ post4 ┆ [2, 5, 3]  ┆ {2,"name2"} │
# │ title5 ┆ post5 ┆ [8, 18, 3] ┆ {3,"name3"} │
# └────────┴───────┴────────────┴─────────────┘

# もう一度実行すればすべての要素をデータフレームに展開できる
df = df.unnest("user")
# ┌────────┬───────┬────────────┬─────┬───────┐
# │ title  ┆ post  ┆ attrs      ┆ id  ┆ name  │
# │ ---    ┆ ---   ┆ ---        ┆ --- ┆ ---   │
# │ str    ┆ str   ┆ list[i64]  ┆ i64 ┆ str   │
# ╞════════╪═══════╪════════════╪═════╪═══════╡
# │ title1 ┆ post1 ┆ [3, 1, 2]  ┆ 1   ┆ name1 │
# │ title4 ┆ post4 ┆ [2, 5, 3]  ┆ 2   ┆ name2 │
# │ title5 ┆ post5 ┆ [8, 18, 3] ┆ 3   ┆ name3 │
# └────────┴───────┴────────────┴─────┴───────┘

# Structの中のリストを展開する（前述したリストのto_sturctを使えばよい）
# attrsと要素が便宜的にリストなだけで、[シェア数, いいね数, ブックマーク数]だとすると
# こうすれば無理やり展開できる
df.with_columns(
    pl.col("attrs").list.to_struct().struct.rename_fields(
        ["share", "favorite", "bookmark"]
    )
).unnest("attrs")
# ┌────────┬───────┬───────┬──────────┬──────────┬─────┬───────┐
# │ title  ┆ post  ┆ share ┆ favorite ┆ bookmark ┆ id  ┆ name  │
# │ ---    ┆ ---   ┆ ---   ┆ ---      ┆ ---      ┆ --- ┆ ---   │
# │ str    ┆ str   ┆ i64   ┆ i64      ┆ i64      ┆ i64 ┆ str   │
# ╞════════╪═══════╪═══════╪══════════╪══════════╪═════╪═══════╡
# │ title1 ┆ post1 ┆ 3     ┆ 1        ┆ 2        ┆ 1   ┆ name1 │
# │ title4 ┆ post4 ┆ 2     ┆ 5        ┆ 3        ┆ 2   ┆ name2 │
# │ title5 ┆ post5 ┆ 8     ┆ 18       ┆ 3        ┆ 3   ┆ name3 │
# └────────┴───────┴───────┴──────────┴──────────┴─────┴───────┘

(余談) OpenAIのBatch APIの出力をパースするときにこれがちょっと便利

OpenAIのBatch APIでGPT-4oなどを使うとこんな感じのjsonlが得られる（便宜的にjqで改行フォーマットしている）と思います。
この例では「JSON形式で返してください」という指示をしています。

{
  "id": "OpenAI側で振られるID",
  "custom_id": "Batch APIの入力で設定した識別子",
  "response": {
    "status_code": 200,
    "request_id": "OpenAI側で振られるRequestのID",
    "body": {
      "id": "OpenAI側で振られるcompletionのID",
      "object": "chat.completion",
      "created": "TIMESTAMP",
      "model": "gpt-4o-2024-05-13",
      "choices": [
        {
          "index": 0,
          "message": {
            "role": "assistant",
            "content": "```json\n{\n    \"Sentiment\": \"negative\",\n    \"Politics topic\": \"no\"\n}\n```",
            "refusal": null
          },
          "logprobs": null,
          "finish_reason": "stop"
        }
      ],
      "usage": {
        "prompt_tokens": "トークン数",
        "completion_tokens": "生成トークン数",
        "total_tokens": "全トークン",
        "prompt_tokens_details": {
          "cached_tokens": 0
        },
        "completion_tokens_details": {
          "reasoning_tokens": 0
        }
      },
      "system_fingerprint": "OpenAI側で振られるフィンガープリント"
    }
  },
  "error": null
}
{
  "id": "OpenAI側で振られるID",
  "custom_id": "Batch APIの入力で設定した識別子",
  "response": {
    "status_code": 200,
    "request_id": "OpenAI側で振られるRequestのID",
    "body": {
      "id": "OpenAI側で振られるcompletionのID",
      "object": "chat.completion",
      "created": "TIMESTAMP",
      "model": "gpt-4o-2024-05-13",
      "choices": [
        {
          "index": 0,
          "message": {
            "role": "assistant",
            "content": "```json\n{\n    \"Sentiment\": \"negative\",\n\"Politics topic\": \"no\"\n}\n```",
            "refusal": null
          },
          "logprobs": null,
          "finish_reason": "stop"
        }
      ],
      "usage": {
        "prompt_tokens": "トークン数",
        "completion_tokens": "生成トークン数",
        "total_tokens": "全トークン数",
        "prompt_tokens_details": {
          "cached_tokens": 0
        },
        "completion_tokens_details": {
          "reasoning_tokens": 0
        }
      },
      "system_fingerprint": "OpenAI側で振られるフィンガープリント"
    }
  },
  "error": null
}

上記のjsonlがopenai_batch_response.jsonlに格納されているとします。

df = pl.read_ndjson("openai_batch_response.jsonl")
# こんな感じで読み込まれます。
# ┌──────────────────────┬─────────────────────────────────┬─────────────────────────────────┬───────┐
# │ id                   ┆ custom_id                       ┆ response                        ┆ error │
# │ ---                  ┆ ---                             ┆ ---                             ┆ ---   │
# │ str                  ┆ str                             ┆ struct[3]                       ┆ null  │
# ╞══════════════════════╪═════════════════════════════════╪═════════════════════════════════╪═══════╡
# │ OpenAI側で振られるID   ┆ Batch APIの入力で設定した識別子     ┆ {200,"OpenAI側で振られるRequest   ┆ null  │
# │                      ┆                                 ┆ のID",…                         ┆       │
# │ OpenAI側で振られるID   ┆ Batch APIの入力で設定した識別子     ┆ {200,"OpenAI側で振られるRequest   ┆ null  │
# │                      ┆                                 ┆ のID",…                         ┆       │
# └──────────────────────┴─────────────────────────────────┴─────────────────────────────────┴───────┘

# まずはレスポンスの中身をカラムに変換
df = df.unnest("response")
# ┌──────────────────────┬─────────────────────────────────┬─────────────┬───────────────────────────────┬─────────────────────────────────┬───────┐
# │ id                   ┆ custom_id                       ┆ status_code ┆ request_id                    ┆ body                            ┆ error │
# │ ---                  ┆ ---                             ┆ ---         ┆ ---                           ┆ ---                             ┆ ---   │
# │ str                  ┆ str                             ┆ i64         ┆ str                           ┆ struct[7]                       ┆ null  │
# ╞══════════════════════╪═════════════════════════════════╪═════════════╪═══════════════════════════════╪═════════════════════════════════╪═══════╡
# │ OpenAI側で振られるID   ┆ Batch APIの入力で設定した識別子     ┆ 200         ┆ OpenAI側で振られるRequestのID   ┆ {"OpenAI側で振られるcompletion    ┆ null  │
# │                      ┆                                 ┆             ┆                               ┆ のID","…                        ┆       │
# │ OpenAI側で振られるID   ┆ Batch APIの入力で設定した識別子     ┆ 200         ┆ OpenAI側で振られるRequestのID   ┆ {"OpenAI側で振られるcompletion    ┆ null  │
# │                      ┆                                 ┆             ┆                               ┆ のID","…                        ┆       │
# └──────────────────────┴─────────────────────────────────┴─────────────┴───────────────────────────────┴─────────────────────────────────┴───────┘

# GPTのresponseの中身から生成テキスト部のみ取り出す
df = df.select(
    pl.col("custom_id"),
    pl.col("body").struct.field("choices").list[0].struct.field("message")),
)
# ┌─────────────────────────────────┬───────────────────────┐
# │ custom_id                       ┆ message               │
# │ ---                             ┆ ---                   │
# │ str                             ┆ struct[3]             │
# ╞═════════════════════════════════╪═══════════════════════╡
# │ Batch APIの入力で設定した識別子     ┆ {"assistant","```json │
# │                                 ┆ {                     │
# │                                 ┆  "Sent…               │
# │ Batch APIの入力で設定した識別子     ┆ {"assistant","```json │
# │                                 ┆ {                     │
# │                                 ┆  "Sent…               │
# └─────────────────────────────────┴───────────────────────┘

# 生成テキスト部を展開
df = df.unnest("message")
# ┌─────────────────────────────────┬───────────┬───────────────────────┬─────────┐
# │ custom_id                       ┆ role      ┆ content               ┆ refusal │
# │ ---                             ┆ ---       ┆ ---                   ┆ ---     │
# │ str                             ┆ str       ┆ str                   ┆ null    │
# ╞═════════════════════════════════╪═══════════╪═══════════════════════╪═════════╡
# │ Batch APIの入力で設定した識別子     ┆ assistant ┆ ```json               ┆ null    │
# │                                 ┆           ┆ {                     ┆         │
# │                                 ┆           ┆  "Sentiment": "negat… ┆         │
# │ Batch APIの入力で設定した識別子     ┆ assistant ┆ ```json               ┆ null    │
# │                                 ┆           ┆ {                     ┆         │
# │                                 ┆           ┆  "Sentiment": "posit… ┆         │
# └─────────────────────────────────┴───────────┴───────────────────────┴─────────┘

# GPT-4oのjsonとして出力させた生成テキストをparseする
df.select(
    pl.col("custom_id"),
    # 接頭の```jsonと末尾```を除いて、json_decodeを実行
    pl.col("content").str.strip_chars("`").str.strip_prefix("json").str.json_decode(),
)
# ┌─────────────────────────────────┬───────────────────┐
# │ custom_id                       ┆ content           │
# │ ---                             ┆ ---               │
# │ str                             ┆ struct[2]         │
# ╞═════════════════════════════════╪═══════════════════╡
# │ Batch APIの入力で設定した識別子     ┆ {"negative","no"} │
# │ Batch APIの入力で設定した識別子     ┆ {"positive","no"} │
# └─────────────────────────────────┴───────────────────┘

# これでGPTの出力を絡むに変換できる
# JSON形式で出力してなくても、フォーマットさえあれば正規表現などを使って抽出は可能
df = df.unnest("content")
# ┌─────────────────────────────────┬───────────┬────────────────┐
# │ custom_id                       ┆ Sentiment ┆ Politics topic │
# │ ---                             ┆ ---       ┆ ---            │
# │ str                             ┆ str       ┆ str            │
# ╞═════════════════════════════════╪═══════════╪════════════════╡
# │ Batch APIの入力で設定した識別子     ┆ negative  ┆ no             │
# │ Batch APIの入力で設定した識別子     ┆ positive  ┆ no             │
# └─────────────────────────────────┴───────────┴────────────────┘

その他便利関数群

特定のgroupごとに要素に対して何らかの計算をしたい時

例えば、group（ex. 学校のクラス）ごとにテストの得点のランキングをつけたいとか、そういう場合。
素直に考えれば、group_byを実行して、その結果を使って要素に対して適用するということが考えられそうです。あるいはSQLでOVER句を使っちゃうということも考えられるかもしれないですね。
ところが、Polarsにもoverという関数があり、SQLのOVER句と同様のことができ、group_byを使う必要がなかったりします。

df.with_columns(
    # クラスごとに得点のランクを付ける
    pl.col("score").rank(descending=True).over("class").alias("score_by_class"),
    # 点数からクラスごとの平均を引く
    (pl.col("score") - pl.col("score").mean().over("class")),
)

複数のデータフレームを結合する

詳しくはpolars.concat — Polars documentationを参照。

pl.concat([df1, df2, df3]) # df1とdf2、df3を縦に結合する
df1.vstack(df2) # df1にdf2を縦に結合する
pl.concat([df1, df2], how="horizontal") # df1とdf2を横方向に結合する

Pythonで定義した関数を実行したい

Pandasでいうapplyみたいな関数が使いたい場合、Polarsではmap_elementsやmap_batchesがそれに当たります。
これを使う場合、RustではなくPythonで動くため、遅くなるという覚悟は必要。

個人的な利用用途としては形態素解析？
MeCabなどを使うときにはtokenize関数を定義して与えたりしています。

def tokenize(text):
    return some_tokenizer.tokenize(text)

def tokenize_batches(texts):
    # HuggingFaceのtokenizerの場合はこれでよさそう
    return some_tokenizer.tokenize(texts)["input_ids"]

pl.with_columns(
    # カラムの要素を一つ一つ関数に与える
    pl.col("text").map_elements(
        tokenize,
        return_dtype=pl.List(pl.String), # 戻り値の型を定義しないとwarningが出る
    ),
    # カラムの要素をまとめて与える方法もあり、並列処理を実装していると便利
    pl.col("text").map_batches(
        tokenize_batches(texts),
        return_dtype=pl.List(pl.Int64),
    ).alias("input_ids"),
)

データフレームの配置し直し

ある項目をカラムにしてデータフレーム配置しなおす（ピボット）

Polarsにはpivot関数がある

# こんなデータフレームを対象とする
# ┌──────┬─────────┬───────┐
# │ name ┆ subject ┆ score │
# │ ---  ┆ ---     ┆ ---   │
# │ str  ┆ str     ┆ i64   │
# ╞══════╪═════════╪═══════╡
# │ A    ┆ math    ┆ 63    │
# │ A    ┆ history ┆ 68    │
# │ B    ┆ math    ┆ 72    │
# │ B    ┆ history ┆ 52    │
# │ C    ┆ math    ┆ 59    │
# │ C    ┆ history ┆ 80    │
# └──────┴─────────┴───────┘

df = df.pivot(
    on="subject",   # どのカラムの値をカラムにするのか
    index="name",   # キー的なカラム
    values="score", # カラムの値としてどれを使うのか
)
# ┌──────┬──────┬─────────┐
# │ name ┆ math ┆ history │
# │ ---  ┆ ---  ┆ ---     │
# │ str  ┆ i64  ┆ i64     │
# ╞══════╪══════╪═════════╡
# │ A    ┆ 63   ┆ 68      │
# │ B    ┆ 72   ┆ 52      │
# │ C    ┆ 59   ┆ 80      │
# └──────┴──────┴─────────┘

ピボットしたテーブルを戻す

Pandasでいうmelt

# 先ほどpivotしたデータフレームを使う
df.unpivot(
    on=["math", "history"],  # 一つのカラムとしてまとめたいカラム名
    index="name",            # キー的なカラム
    variable_name="subject", # 現在のカラム名をまとめたカラムにつける名前
    value_name="score",      # カラムの中の値をまとめたカラムにつける名前
)
# ┌──────┬─────────┬───────┐
# │ name ┆ subject ┆ score │
# │ ---  ┆ ---     ┆ ---   │
# │ str  ┆ str     ┆ i64   │
# ╞══════╪═════════╪═══════╡
# │ A    ┆ math    ┆ 63    │
# │ B    ┆ math    ┆ 72    │
# │ C    ┆ math    ┆ 59    │
# │ A    ┆ history ┆ 68    │
# │ B    ┆ history ┆ 52    │
# │ C    ┆ history ┆ 80    │
# └──────┴─────────┴───────┘

以上です。
紹介したのはPolarsの機能のうちほんの一部の自分がよく使う機能なので、皆さんの用途に適切な機能がもっと他にもあると思います。
Pandasを普段使っているような人であれば、Polarsはすぐに使えるようになるので、ぜひ使ってみてください。
ぜひ、公式ドキュメントを漁って探してみてください。

明日のアドベントカレンダーは松村さんの「SolanaでSBT（Soulbound Token）を発行してみた」です。