Pandas の代わりに Polars でデータ加工する豆知識

注意

かゆいところに手が届くようにするためのメモ。

.select / .filter 関連

列の選択に関する事柄。

pl.col()

Polars における「列」を表現するクラスを生成できる。pl.col() すると割と好き放題でき、列に対して四則演算したり文字列扱いして加工処理を回したりできる。

また、str な列については pl.col().str すると列を str な値として扱え、contains() とか ends_with() とかいかにも便利そうな関数を使えるようになる。

以下は id 列が偶数かつ email 列に .com を含む filter 処理を回す例である。

df = (
    pl.read_csv(f"{dirname(__file__)}/vendor/MOCK_DATA.csv")
    .filter(pl.col("id") % 2 == 0)
    .filter(pl.col("email").str.contains(".com"))
    .select([pl.col("id"), pl.col("email")])
)

print(df)

出力：

shape: (307, 2)
┌─────┬────────────────────────────┐
│ id  ┆ email                      │
│ --- ┆ ---                        │
│ i64 ┆ str                        │
╞═════╪════════════════════════════╡
│ 2   ┆ [e-mail]       │
│ 4   ┆ [e-mail]        │
│ 6   ┆ [e-mail]      │
│ 10  ┆ [e-mail]       │
│ 12  ┆ [e-mail] │
│ …   ┆ …                          │
│ 982 ┆ [e-mail]       │
│ 984 ┆ [e-mail]    │
│ 992 ┆ [e-mail]      │
│ 996 ┆ [e-mail]    │
│ 998 ┆ [e-mail]     │
└─────┴────────────────────────────┘

pl.exclude()

指定した列「以外」の列を表現するクラスを生成できる。DataFrame.filter() と組み合わせると便利である。

例：

import polars as pl

data = {"alpha": [1, 2, 3], "beta": [1, 2, 3], "gamma": [1, 2, 3]}

df = pl.DataFrame(data).select([pl.exclude("alpha")])

print(df)

出力：

shape: (3, 2)
┌──────┬───────┐
│ beta ┆ gamma │
│ ---  ┆ ---   │
│ i64  ┆ i64   │
╞══════╪═══════╡
│ 1    ┆ 1     │
│ 2    ┆ 2     │
│ 3    ┆ 3     │
└──────┴───────┘

列の追加 (with_columns)

Pandas だと df["col_1"] = ... みたいな感じで列を生やせたが、Polars では DataFrame.with_columns() で生やすことになる。

このメソッドには罠っぽいところが一つあり、with_columns は破壊的なメソッドではない。つまり、with_columns したデータフレームを直接更新するわけでないから、必要であれば結果をデータフレームに再代入する必要がある。

以下は email 列のドメイン部分を取り出して(@ で文字列を分割し、2つめの要素をドメインとする)、それを email_domain 列としてデータフレームに追加する例である。

df = pl.read_csv(f"{dirname(__file__)}/vendor/MOCK_DATA.csv")
df = df.with_columns(df["email"].str.split("@").list[1].alias("email_domain"))

print(df.select([pl.col("id"), pl.col("email"), pl.col("email_domain")]))

LazyFrame 関連

以下のような感じで CSV ファイルを読み込むものとする。CSVの列はコードに書かれていないものも含めて7列ある。

from posixpath import dirname
import polars as pl

df = (
    pl.scan_csv(f"{dirname(__file__)}/vendor/MOCK_DATA.csv")
    .select(
        [
            pl.col("first_name"),
            pl.col("city"),
            pl.col("email"),
            pl.col("ip_address"),
        ]
    )
    .filter(pl.col("email").str.ends_with(".com"))
    .select(pl.col("email"))
)

pl.scan_csv() / LazyFrame.collect()

pl.read_csv と似ているが、pl.scan_csv() は LazyFrame.collect() が実行されるまで実際の読み込みを待機する。これによって後続のクエリの指定を待ったうえで CSV を読み込めるため、読み込み処理のメモリ効率・処理効率の向上が見込める。

# .collect() すると LazyFrame から DataFrame になる
pl.scan_csv("file.csv").collect()

LazyFrame.explain()

LazyFrame に与えられたクエリに基づいたクエリプランを提示する。optimized プロパティ経由で最適化された・されていないクエリプランを出せる。

以下を例にすると、最適化前は7列すべてを読み込んでから二回 SELECT しているが、最適化された後では読み込みの段階で1列だけに絞ったうえで一回だけ SELECT するので済んでいる。

print(df.explain(optimized=False))
print(df.explain(optimized=True))

 SELECT [col("email")] FROM
  FILTER col("email").str.ends_with([String(.com)]) FROM

   SELECT [col("first_name"), col("city"), col("email"), col("ip_address")] FROM

      Csv SCAN /home/ubuntu/projects/python-sandbox/src/vendor/MOCK_DATA.csv
      PROJECT */7 COLUMNS
FAST_PROJECT: [email]

    Csv SCAN /home/ubuntu/projects/python-sandbox/src/vendor/MOCK_DATA.csv
    PROJECT 1/7 COLUMNS
    SELECTION: col("email").str.ends_with([String(.com)])

LazyFrame.show_graph()

GraphViz を使用して、explain() の出力をグラフとして出す。

df.show_graph(optimized=True, show=False, output_path="optimized.png")
df.show_graph(optimized=False, show=False, output_path="not_optimized.png")

optimized=True:

optimized=False: