https://www.m3tech.blog/entry/python/pavise

AI・機械学習チームの北川(@kitagry)です。この記事はAI・機械学習チームブログリレー4日目の記事です。前日は横本(@yokomotod)さんの『distrolessコンテナイメージの中を覗いて「なんか軽くてセキュアらしい」より理解を深める』でした。

www.m3tech.blog

pandasやpolarsを利用していると、「このDataFrameは何のカラムを持っているんだ？」と悩まされたことはないでしょうか。型ヒントをつけたくても pd.DataFrame や pl.DataFrame という型では列の情報を表現できず、説明をきちんと書かなかった過去の自分を恨むということは誰しもが通った道ではないでしょうか。

このような課題に対してはまず pandera が有力な選択肢です。実際にAI・機械学習チームでもpanderaを導入していました。しかし使い込んでいくうちに、特定のシナリオで苦しいポイントが見えてきました。

www.m3tech.blog

この記事では、panderaで実際に詰まったポイントと、その課題を解決するために開発した pavise を紹介します。

DataFrameを型安全に扱いたい

pandas や polars のDataFrameは非常に柔軟ですが、その柔軟さゆえに型情報が失われがちです。

# pd.DataFrame だと次のように列情報がコードに現れない
def process_users(df: pd.DataFrame) -> pd.DataFrame:
    return df[df["age"] >= 18]  # "age" が存在するか、型は何かがコードからわからない

列名のタイポや型の不一致は実行時まで発覚せず、大規模なMLパイプラインでは原因特定も難しくなります。

panderaで解決できること

pandera を使うと、DataFrameのスキーマを定義して実行時にバリデーションできます。

import pandera as pa
from pandera.typing import DataFrame, Series

class UserSchema(pa.DataFrameModel):
    user_id: Series[int]
    name: Series[str]
    age: Series[int] = pa.Field(ge=0, le=150)

def process_users(df: DataFrame[UserSchema]) -> DataFrame[UserSchema]:
    return df[df["age"] >= 18]

スキーマを定義することで、列名・型・バリデーションルールをコードで表現でき、実行時に違反があればエラーを返してくれます。DataFrameに型安全を導入したいなら、まずpanderaを検討することをおすすめします。

panderaで苦しかった2つのポイント

しかし、大規模なシステムで使い込んでいくと、2つの点で苦しさを感じました。

1. テストで使わないカラムを全部埋めなければならない

DataFrame[UserSchema] という型を使うと、テスト時もSchemaを満たすDataFrameを渡す必要があります。しかし実際には、テスト対象の関数が使うカラムは一部だけのことがほとんどです。

def filter_adults(df: DataFrame[UserSchema]) -> DataFrame[UserSchema]:
    return df[df["age"] >= 18]  # 使うのは "age" だけ

def test_filter_adults():
    df = pd.DataFrame({
        "user_id": [1, 2, 3],
        "name": ["Alice", "Bob", "Charlie"],
        "age": [15, 20, 25],
        "email": ["a@example.com", "b@example.com", "c@example.com"],
        "registered_at": [datetime(2020, 1, 1)] * 3,
        "segment_a": [False, False, False],
        "segment_b": [False, False, False],
        "segment_c": [False, False, False],
        # ... ageだけテストしたいのにSchemaの全カラムを埋めている
    })
    result = filter_adults(DataFrame[UserSchema](df))
    assert len(result) == 2

Schemaのカラム数が増えるほど、テストデータの準備が大変になっていきました。

2. 依存の方向が逆になる

データパイプラインでは「データソースからDataFrameを作る部分」と「そのDataFrameを使ってモデルを学習・推論する部分」を分離することがよくあります。このとき依存の方向として自然なのは次の構造です。

モデル部分（feature_a, feature_b が必要）← データ取得部分（それを含むDataFrameを返す）

しかし pandera のSchemaはクラス継承ベースであるため、両者で同じSchemaクラスを共有するか、継承関係を作る必要があります。

# pandera: データ取得部分のSchemaをモデル部分が参照してしまう
class RawDataSchema(pa.DataFrameModel):  # データ取得部分が定義
    user_id: Series[int]
    feature_a: Series[float]
    feature_b: Series[float]
    # ... BQ由来の大量のカラム

# モデル部分は feature_a, feature_b だけ使いたいのに RawDataSchema に依存してしまう
def train_model(df: DataFrame[RawDataSchema]) -> Model:
    X = df[["feature_a", "feature_b"]].to_numpy()
    ...

モデル部分が本来知らなくていいカラムの情報まで抱えることになり、依存の方向が逆になっていました。

paviseのアプローチ——Protocolで解決

pavise は Python 標準の Protocol を使ってDataFrameのスキーマを定義します。

from typing import Protocol
from pavise.pandas import DataFrame

class AgeFilterable(Protocol):
    age: int

def filter_adults(df: DataFrame[AgeFilterable]) -> DataFrame[AgeFilterable]:
    return df[df["age"] >= 18]

Protocol の特性により、「int型のage カラムを持つ任意のDataFrame」を受け付けることができます。加工層は自分が必要なカラムだけを定義すればよく、ダウンロード層がどんなSchemaを持っていても、そのカラムを含んでいれば型チェッカーが通ります。

class FullUserSchema(Protocol):
    user_id: int
    name: str
    age: int
    email: str
    # ... 他のカラム

# FullUserSchema は AgeFilterable を満たすので型チェッカーが通る
full_df: DataFrame[FullUserSchema] = DataFrame[FullUserSchema](raw_df)
filtered = filter_adults(full_df)  # OK

これにより、依存の方向を自然に保てます。そしてテスト時も、テスト対象の関数が必要とするカラムだけを持つDataFrameを渡せばよくなります。

# テストでは "age" だけ持つ最小限のDataFrameを渡せる
def test_filter_adults():
    df = pd.DataFrame({"age": [15, 20, 25]})
    result = filter_adults(DataFrame[AgeFilterable](df))
    assert len(result) == 2

もし Protocol に定義したカラムが存在しない、または型が一致しない場合、mypy/pyright が静的に検出してくれます。実行前に問題を発見できるのは大きなメリットです。

paviseの機能

バリデータで値の制約を表現

Annotated を使って値の制約を型として表現できます。

from typing import Protocol, Annotated
from pavise.validators import Range, Regex

class UserSchema(Protocol):
    age: Annotated[int, Range(0, 150)]
    email: Annotated[str, Regex(r'^[\w.+-]+@[\w-]+\.[a-zA-Z]{2,}$')]

型チェッカーの機能でデータ取得部分にだけバリデーションを行い、モデル部分にはバリデーションがすでに終わっている前提で書くことも可能です。

# データ取得部分
class RawDataSchema(pa.DataFrameModel):
    age: Annotated[int, Range(0, 150)]  # データ取得部分ではバリデーションを通して厳しくする


# モデル部分
class AgeFilterable(Protocol):
    age: int

def filter_adults(df: DataFrame[AgeFilterable]) -> DataFrame[AgeFilterable]:
    return df[df["age"] >= 18]

filter_adults(DataFrame[RawDataSchema](df))  # 型チェッカーはOK

`NotRequiredColumn` で存在しないカラムを許容

「カラム自体が存在しないこともある」と「値が None のこともある」を区別できます。

from typing import Optional
from pavise.pandas import DataFrame, NotRequiredColumn

class UserSchema(Protocol):
    user_id: int
    age: int | None                           # 列は必須、値はNone許容
    email: NotRequiredColumn[Optional[str]]   # 列自体がなくてもよい

まとめ

DataFrameへの型導入としてはまず pandera が実績もあるため有力な選択肢です。スキーマ定義・実行時バリデーションを手軽に始められます。

一方で、大規模なMLシステムで依存の方向を整理したい場合や、テスト時に不要なカラムの準備が煩雑になってきた場合は、Protocol ベースの pavise が助けになるかもしれません。

panderaでもpaviseでも、DataFrameに型を導入することで「実行してみないとわからない」エラーを静的解析や境界でのバリデーションで事前に潰せます。大規模なMLパイプラインを安全に育てていくために、ぜひ型安全なDataFrameの運用を検討してみてください。

Pavise Documentation — Pavise documentation

pip install pavise[pandas]
# または
pip install pavise[polars]

We are Hiring!

弊社ではPythonの型によって安全なプロダクトを作成できるようなエンジニアを募集しています。興味がある方は次のリンクから応募をお待ちしています。