Polarsのlazy APIで速くなる仕組み—pandasから書き換えて実測

Polarsのlazy APIで速くなる仕組み—pandasから書き換えて実測 | mohablog

Polarsは「pandasより速いデータフレーム」として広まりました。ただ、pandasも3.0でCopy-on-Writeが標準になり内部最適化が進んだことで、単純な集計だと差はそれほど開きません。本当に効くのはlazy APIのクエリ最適化のほうです。

この記事はPolars 1.41.0pandas 3.0.3、Python 3.13で実測しながら、どこで差が出てどこで出ないかを切り分けます。

目次

Polarsを最初の数行で動かす

CSVを読み、列で絞り、グループ集計するまでを数行で書けます。pandasを触ったことがあれば、メソッド名の対応はだいたい想像どおりです。

scan_csvとcollectの最小例

まず目を引くのがscan_csvcollectの組み合わせ。読み込みを宣言した時点では何も走りません。

import polars as pl

q = (
    pl.scan_csv("sales.csv")        # この行ではまだCSVを読まない
    .filter(pl.col("value") > 500)
    .group_by("category")
    .agg(total=pl.col("value").sum())
    .sort("category")
)

print(q.collect().head(3))          # collect() で初めて実行される
shape: (3, 2)
┌──────────┬──────────┐
│ category ┆ total    │
│ ---      ┆ ---      │
│ i64      ┆ f64      │
╞══════════╪══════════╡
│ 0        ┆ 1.5013e8 │
│ 1        ┆ 1.4974e8 │
│ 2        ┆ 1.4971e8 │
└──────────┴──────────┘

eagerとlazyの2つの入口

Polarsには入口が2つあります。read_csvはその場で読み込むeagerscan_csvは実行を遅らせるlazy。lazyは式を組み立てるだけで、collect()を呼ぶまで一切計算しません。後述のクエリ最適化はこのlazy側でしか効きません。

速さの出どころと、pandas 3.0で縮んだ差

列指向とマルチスレッドという土台

Polarsは列指向でデータを保持し、デフォルトで全コアを使います。pandasのDataFrame演算は基本シングルスレッド。これが素の速度差の土台です。ただ、土台の差がそのまま実行時間の差になるとは限りません。

単純なgroupbyだと差は1.4倍

2000万行(CSVで517MB)をメモリに載せた状態で、value > 500で絞ってcategory別に合計と平均を出す。この程度の集計だと実測はこうなりました。

pandas eager   : 0.169s  rows=50
polars eager   : 0.121s  rows=50
speedup        : 1.4x

pandas 3.0はCopy-on-Writeが標準になり、不要なコピーが減りました。文字列はPyArrowバックエンドで5〜10倍速くなりますが、今回は数値だけの集計なので効きません。数値のPyArrow化はdtype_backend="pyarrow"での任意適用のまま。それでも素の集計でPolarsに肉薄します。「常に10倍速い」という話は、古いpandasや重い処理が前提のことが多い。pandasの基本操作と同じ感覚で書けて、この集計では速度もほとんど変わりません。差が大きく開くのは、ディスクから読み込む処理を挟んだときです。

lazy APIのクエリ最適化で読む量を減らす

公式ドキュメントの “Lazy API” は、遅延実行の利点をこう締めています。「the lazy API should be preferred unless you are either interested in the intermediate results or are doing exploratory work」。途中結果が要るか試行錯誤中でなければlazyを使え、という立場です。理由は、Polarsがクエリ全体を見てから実行計画を組み直すから。

explain()で最適化の前後を見る

lazyクエリはexplain()で実行計画を覗けます。最適化なし(optimized=False)と、ありの計画を並べると違いが一目で分かります。

q = (
    pl.scan_csv("sales.csv")
    .filter(pl.col("value") > 500)
    .group_by("category")
    .agg(total=pl.col("value").sum(), avg_qty=pl.col("qty").mean())
)
print(q.explain(optimized=False))
print(q.explain(optimized=True))
# 出力はalias等を一部省略して掲載
# optimized=False(書いたとおりの計画)
AGGREGATE
  [col("value").sum(), col("qty").mean()] BY [col("category")]
  FROM
  FILTER [(col("value")) > (500.0)]
  FROM
    Csv SCAN [sales.csv]
    PROJECT */4 COLUMNS

# optimized=True(最適化後)
AGGREGATE
  [col("value").sum(), col("qty").mean()] BY [col("category")]
  FROM
  Csv SCAN [sales.csv]
  PROJECT 3/4 COLUMNS
  SELECTION: [(col("value")) > (500.0)]

FILTERという独立した工程が消え、SELECTIONとしてCsv SCANの中に取り込まれました。さらにPROJECT*/4から3/4に。読む列も絞る行も、スキャンの段階に押し下げられています。

predicate pushdown:フィルタがスキャンに降りる

公式 “Lazy API” のPredicate Pushdownはこう説明されます。「Apply filters as early as possible while reading the dataset」。読みながら早い段階でフィルタを適用するという意味。上の計画でSELECTIONがスキャンに入ったのがこれです。全行をメモリに展開してから捨てるのではなく、読む時点で落とす。

projection pushdown:要らない列は読まない

もう一つがProjection Pushdown。「Select only the columns that are needed while reading the dataset」とある通り、必要な列だけ読みます。集計に使うのはcategoryvalueqtyの3列。残り1列(region)はディスクから読まれもしません。計画のPROJECT 3/4 COLUMNSがその証拠です。業務で30列あるログCSVを2列だけ集計したとき、列を手で絞らずにscan_csvへ渡すだけで読み込み量が減ったのは、この最適化が効いていたからでした。pandasのread_csvは基本的に全列を読み込みます。

pandasからの書き換え対応

よく使う操作の対応はこの通り。メソッド名が変わるだけのものが多い。

pandasPolars
df[df["x"] > 5]df.filter(pl.col("x") > 5)
df.groupby("a").agg(...)df.group_by("a").agg(...)
df["a"].mean()df.select(pl.col("a").mean())
df.assign(c=df["a"] + 1)df.with_columns(c=pl.col("a") + 1)
df.merge(other, on="id")df.join(other, on="id")
df.sort_values("x")df.sort("x")

列の追加はwith_columns、列の選択や集計はselectに集約されます。この2つを押さえると大半の処理は書けます。

移行でハマる書き方

pandasの癖をそのまま持ち込むと、速さが出ないどころか動かないコードになります。詰まりやすいのは2か所。

applyを式に書き直していない

pandasのapply感覚でmap_elementsにPython関数を渡すと、行ごとにPythonを呼ぶので遅くなります。

# 遅い:行ごとにPython関数を呼ぶ
df.with_columns(
    taxed=pl.col("value").map_elements(lambda x: x * 1.1, return_dtype=pl.Float64)
)

# 速い:式(expression)でベクトル化
df.with_columns(taxed=pl.col("value") * 1.1)
shape: (5, 3)
┌──────────┬────────────┬────────────┐
│ category ┆ value      ┆ taxed      │
│ ---      ┆ ---        ┆ ---        │
│ i64      ┆ f64        ┆ f64        │
╞══════════╪════════════╪════════════╡
│ 4        ┆ 96.782484  ┆ 106.460732 │
│ 38       ┆ 679.929338 ┆ 747.922271 │
│ 32       ┆ 292.516999 ┆ 321.768699 │
└──────────┴────────────┴────────────┘

式で書けば全コアで並列に処理されます。map_elementsは式で表現できない処理に限った最後の手段。

インデックス前提のコード

Polarsに行インデックスはありません。df.loc[]set_indexは存在しない。行の絞り込みはfilter、特定行の取り出しはrowsliceで書き直します。pandasの「インデックスで結合」を多用していたコードは、joinのキー指定に置き換える必要があります。

メモリに載らないデータはstreaming

データがRAMに収まらないときは、collect(engine="streaming")でストリーミング実行に切り替えます。公式 “Streaming” の説明はこう。「Instead of processing all the data at once, Polars can execute the query in batches allowing you to process datasets that do not fit in memory」。一括ではなくバッチで処理する、というだけの違いです。

同じ517MBのCSVを、読み込みから集計・ソートまで通して測るとこうなりました。scan_csvでディスクから読む経路だと、最適化とストリーミングの効果がそのまま出ます。

pandas (read_csv→filter→groupby→sort): 1.815s   peak 2257MB
polars lazy + streaming               : 0.187s   peak  799MB

実行時間で約9.7倍、ピークメモリは2257MBから799MBへ。in-memoryの1.4倍とはまるで別の数字です。差を生んだのは演算速度より、読む量を減らしたことのほうでした。

まとめ

  • in-memoryの単純な集計では、pandas 3.0との差は1.4倍程度にとどまる
  • scan_*とlazy APIを使うと、predicate/projection pushdownで読む行と列が減る
  • explain(optimized=True)で、最適化後の実行計画を実際に確認できる
  • RAMに載らないデータはcollect(engine="streaming")でバッチ処理に切り替える
  • pandasのapplyやインデックス前提のコードは持ち込まず、式とfilterで書き直す
よかったらシェアしてね!
  • URLをコピーしました!
  • URLをコピーしました!
目次