AI コーディングの時代になり、作ってみたいというぼんやりとしたアイデアだけはあったものを作るハードルが下がっているとよく言われる。その流れに乗じて自分の中にあったアイデアを実装してみることにした。
それはデータフレームの型チェッカーだ。Pandas や Polars などでのデータ処理を書いていると、こういう「実行するまで気づけない」ミスがよくある。
- 集計した結果は
Float64なのに、出力スキーマではintのつもりだった - 結合やグルーピングのキー名を間違えていて、そもそも結合できない列を指定していた
- 上流の集計でカラムが消えていて、後続の処理で参照できない
データ分析系の仕事をしていると、こうしたことが原因で起きる実行時エラーに苛立たせられることは少なからずある。しかもこうしたエラーは往々にして、少し重いデータで走らせてみようかとかいろんなデータで走らせてみようといったタイミングで、長時間実行した後で見つかる。
こうした時に静的なチェックで先に「クサい」部分を洗い出しておければと思う。しかし、 mypy などの既存の型チェッカーでは DataFrame の「中身（列名・型）」まで見られないのでこうしたミスは検出できない。テストやパイプラインを実際に走らせて、データを流して初めてエラーになる。
そんなことにイライラしていた頃にレコード多相性(record polymorphism)の概念に触れた。データが含むレコードまで取り込んだ多相性のある型システムを作ることで型検査の段階で可能な計算式か層でないかを区別できる。これを使えば上に書いたような苛立ちも消えるのではないか！ しかしなかなかデータ分析の仕事を ML でやることもできない。静的型チェックツールにこうした機能をサポートするものがあれば……と思いながら時は過ぎ、AI エージェント全盛時代になってしまった。
意外と AI エージェントもこうしたデータフレームについて正確なコードを書くのには苦労するような体感もあるし、シンプルに作ってみるのは面白そうだなと思ったのである。
このツールはレコード多相性の概念に則ってそのまま実装している訳ではないが、それが与えてくれる恩恵を多少なりと得られる段階までは作ることができたと思う。もし何らかのフィードバックなどがあれば嬉しいと思う。
- リポジトリ: https://github.com/ngr-t/polypolarism
- 対応バージョン
- Polars 1.37 以上
- Pandera 0.19 以上
スコープ
一旦 Polars に対応することとし、型アノテーションの記述は Pandera Schema を用いることにした。
Pandas よりも Polars の API の方が比較的すっきりしているので実装が楽になるであろうという判断はある。
本来のレコード多相性を実装するには Pandera の記述力では不十分なのだが、実用上の多くのケースでは記述力に問題はないし、わざわざ DSL で書くよりは定着しつつあるライブラリに乗っかる方が良いと考えた。実行時の検証ができるのももちろん良い点だ。
Quick Start
インストール
GitHub から直接インストールする。
pip install git+https://github.com/ngr-t/polypolarism.git
ソースから入れる場合:
git clone https://github.com/ngr-t/polypolarism.git
cd polypolarism
pip install .
使い方
スキーマを Pandera の DataFrameModel で宣言し、関数の引数・戻り値に DataFrame[Schema] 型を付けるだけだ。
import polars as pl
import pandera.polars as pa
from pandera.typing.polars import DataFrame
class Users(pa.DataFrameModel):
user_id: int
name: str
class Orders(pa.DataFrameModel):
order_id: int
user_id: int
amount: pl.Float64
class Joined(pa.DataFrameModel):
user_id: int
name: str
order_id: int
amount: pl.Float64
def merge_users_orders(
users: DataFrame[Users],
orders: DataFrame[Orders],
) -> DataFrame[Joined]:
return users.join(orders, on="user_id", how="inner")
あとはファイルを指定して実行するだけだ。
polypolarism your_module.py
問題なければこう表示される。
merge_users_orders: OK
aggregate_sales: OK
All 2 function(s) passed.
型エラーを静的に検出する例
ここがこのツールの一番おいしいところだ。次は join → 列の派生 → group_by().agg() という、よくある多段パイプライン。sum() の結果が Float64 になるのに、出力スキーマ RevenueByRegion.total_revenue を int（Int64）で宣言してしまっている。
class Sales(pa.DataFrameModel):
order_id: int
region: str
product_id: int
quantity: int
class Products(pa.DataFrameModel):
product_id: int
unit_price: pl.Float64
class RevenueByRegion(pa.DataFrameModel):
region: str
total_revenue: int # Int64 と宣言しているが、sum(Float64) は Float64 になる
def compute_revenue(
sales: DataFrame[Sales],
products: DataFrame[Products],
) -> DataFrame[RevenueByRegion]:
return (
sales
.join(products, on="product_id", how="inner")
.with_columns(revenue=pl.col("quantity") * pl.col("unit_price"))
.group_by("region")
.agg(total_revenue=pl.col("revenue").sum())
)
このコードを実行せずにチェックすると、ちゃんとエラーになる。
compute_revenue: FAIL
- [PLY040] Column 'total_revenue' has type Float64, but declared type is Int64
1 function(s) failed, 0 passed.
join のキー間違いなど、別の種類のエラーも診断コード付きで報告される。
bad_join: FAIL
- [PLY010] Column 'user_id' not found in right frame
- [PLY040] Could not infer return type
1 function(s) failed, 0 passed.
validation-as-narrowing：素性のわからない DataFrame に型を効かせる
スキーマで検証した時点から型を絞り込んでくれるようにもなっている。次のいずれの書き方でも narrowing が効く。
Schema.validate(df)df.pipe(Schema.validate)Schema.validate(lf).collect()
たとえば、外から pl.DataFrame（スキーマ注釈のない、素性のわからない DataFrame）として渡ってきても、一度 validate を通せば、そこから先は「そのスキーマの DataFrame」として静的に追ってくれる。
class Raw(pa.DataFrameModel):
user_id: int
amount: pl.Float64
class Config:
strict = True
class TotalByUser(pa.DataFrameModel):
user_id: int
total: pl.Float64
class Config:
strict = True
@pa.check_types
def total_ok(df: pl.DataFrame) -> DataFrame[TotalByUser]:
df = Raw.validate(df) # ここから df は DataFrame[Raw] として追跡される
return df.group_by("user_id").agg(total=pl.col("amount").sum())
@pa.check_types
def total_bug(df: pl.DataFrame) -> DataFrame[TotalByUser]:
df = Raw.validate(df)
return df.group_by("user_id").agg(total=pl.col("amont").sum()) # タイポ: amount -> amont
total_bug は集計列を amount のつもりで amont とタイプミスしている。validate で Raw に絞り込まれているので、その後の集計でも列の取り違えを拾える（Raw を strict にしておくと、宣言にない列への参照をこうして弾ける）。
total_ok: OK
total_bug: FAIL
- [PLY011] Aggregation column 'amont' not found in frame from schema 'Raw'
- [PLY040] Could not infer return type
1 function(s) failed, 1 passed.
エラーメッセージが from schema 'Raw' と言っているとおり、validate 以降は Raw として扱われているのが分かる。素性のわからない DataFrame でも、入口で一度通しておけば、そこから先は型が効くようになる、という使い方だ。
分岐があっても追ってくれる
if などの制御構文で出力の作り方が変わるような関数でも、分岐をすべて追って、合流させた上でチェックしてくれる。
次の 2 つは同じ形をしているが、add_a_bug のほうは else 側だけ列 a のつもりで b を作ってしまっている。
class KV(pa.DataFrameModel):
k: str
v: pl.Float64
class Config:
strict = True
class KVa(pa.DataFrameModel):
k: str
v: pl.Float64
a: pl.Float64
class Config:
strict = True
@pa.check_types
def add_a_ok(df: DataFrame[KV], flag: bool) -> DataFrame[KVa]:
if flag:
out = df.with_columns(a=pl.col("v") * 2)
else:
out = df.with_columns(a=pl.col("v") + 1)
return out
@pa.check_types
def add_a_bug(df: DataFrame[KV], flag: bool) -> DataFrame[KVa]:
if flag:
out = df.with_columns(a=pl.col("v") * 2)
else:
out = df.with_columns(b=pl.col("v") + 1) # ミス: a のつもりが b
return out
両方の分岐を見た上で、add_a_bug の else 側で a が作られないこと（と、要らない b ができていること）を指摘してくれる。
add_a_ok: OK
add_a_bug: FAIL
- [PLY040] Missing column 'a' of type Float64
- [PLY040] Extra column 'b' of type Float64 not in declared return type
1 function(s) failed, 1 passed.
if/else だけでなく、三項式（A if cond else B）や for ループ内の早期 return、elif の途中の分岐に同じミスを仕込んでも、手元（0.1.0）ではいずれも同じように検出できた。
CLI
ファイル単体でもディレクトリ単位でもチェックできる。エディタ連携用に JSON 出力も用意してみた。vscode 連携も作っている途中である。
# 単一ファイル or ディレクトリをチェック
polypolarism path/to/file.py
polypolarism path/to/project/
# JSON 出力（エディタ連携など）
polypolarism --format json path/to/file.py
# その他のフラグ
polypolarism --version
polypolarism --no-color path/to/file.py
polypolarism --polars-version 1.40 path/to/file.py # PLW010 を抑制
polypolarism --no-version-check path/to/file.py