KAKEHASHI Tech Blog

カケハシのEngineer Teamによるブログです。

データサイエンスでもっと型ヒントの恩恵を受けるための簡単なpandas拡張

こんにちは。Musubi AI 在庫管理の開発チームで機械学習エンジニアをしている保坂です。

こちらの記事はカケハシ Advent Calendar 2022 の 15 日目の記事になります。

今日はMusubi AI在庫管理の需要予測や発注おすすめ作成といったデータサイエンスロジックの開発において活用している、 型ヒントの恩恵を受けやすくするための簡単なpandas拡張についてご紹介したいと思います。

Pythonはバージョン3.5より型ヒントがサポートされるようになり、静的型付け言語と同等とまでは行かないまでも、同様の恩恵が受けられるようになってきましたね。 Pythonにおける型ヒントの恩恵としては以下のようなものがあると思います。

  • mypyによる静的型チェック
  • pydanticによる動的型チェック
  • IDE/エディタによる入力候補の表示
  • 可読性向上

しかし依然として、データサイエンスでよく使われるpandasDataFrame クラスを用いた場合、型ヒントの恩恵を十分に受けられないと感じています。

DataFrame クラスは複数の行や列を持つテーブルデータを1インスタンスで扱う事ができるものです。複数の行や列に対して一括で操作を適用したり、手軽に集計を行ったりする事ができるので、データ分析やデータサイエンスロジックの開発では手放せないものです。

ただ、DataFrame では列に関する定義を型ヒントで考慮することができず、データパイプラインでどういうふうにデータが変換されていくのかがわかりにくいと思うことがよくあります。

例えばデータの変換のために以下のようなプロトタイプ宣言を持つ関数を作成することがありますが、これを見ても変換内容に関してわかることはほとんどありません。

>>> def data_processing(input_data: pd.DataFrame) -> pd.DataFrame:
...     ...

これだけだと、 DataFrame が別の DataFrame に変換されることはわかるが、どんな列を持つテーブルが入力となり、どんな列を持つテーブルが出力されるかがわかりません。 特に、これまでエンジニアと一緒に開発をした際にはまずこの点を指摘されていました。docstringでDataFrameの列について記載するようにしていたこともありますが、コードと乖離しないようにするのが大変だった記憶があります。

作成した DataFrame 拡張

このような課題を解決するため、DataFrame クラスに型ヒントやバリデーション機能を追加した TypedDataFrame クラスを作りました。

車輪の再発明は極力避けるべきことだと思うので、DataFrame でも型ヒントに相当する機能が使えるパッケージはないかと調査し、 いくつかのパッケージを見つけることができました。

しかし利用者数が少ないなど様々な理由から導入せず、自前で簡単なものを作ることにしました。

なお、AI在庫管理の開発がスタートした2020年当時は発見することができなかったのですが、最近はpanderaが広く使われるようになってきていて、利用者が多く、同様の機能を提供しているため、こちらを活用するほうがよいかもしれません。

以下の様な考えのもと、TypedDataFrame が複雑になりすぎないように務めました。

  • 複雑でこみいったコード・難しい仕様を活用したコードにならない範囲でDataFrameを拡張する
  • 将来このような機能を提供するオープンソースソフトウェアが登場した際にはできるだけ乗り換えやすいように、簡潔な仕様に留める

作成した TypedDataFrame は以下のような DataFrame のサブクラスです。 実際に利用しているクラスから最小限の部分を抜粋・改変して紹介しています。 このコードはPython 3.8, pandas 1.3.5 で動作するものです。

>>> from typing import Any, Dict, Type, TypeVar
>>> import pandas as pd
>>> TypedDataFrameSubclass = TypeVar("TypedDataFrameSubclass", bound="TypedDataFrame")

>>> class TypedDataFrame(pd.DataFrame):
...     _dtype: Dict[str, Type] = {"dummy": int}
...     def __init__(self, *args: Any, **kwds: Any) -> None:
...         super().__init__(*args, **kwds)
...         self._validate_columns()
...         self._change_dtypes()
...     def _validate_columns(self) -> None:
...         msg = (
...             f"入力データの列がテーブル定義と違いますよ\n"
...             f"入力データ: {sorted(self.columns)}\n"
...             f"テーブル定義: {sorted(self._dtype.keys())}"
...         )
...         assert sorted(self.columns) == sorted(self._dtype.keys()), msg
...     def _change_dtypes(self) -> None:
...         for col, dtype in self._dtype.items():
...             self[col] = self[col].astype(dtype)
...     @classmethod
...     def read_csv(cls, path: str) -> TypedDataFrameSubclass:
...         return cls(pd.read_csv(path, dtype=cls._dtype))
...

pandasを拡張する方法はExtending pandas に示されており、これを参考にしています。 ただ、このドキュメントに記載されている内容は実施したい拡張の用途とは外れるため特に適用はしていません。

TypedDataFrame の利用方法

このクラスの利用方法をご紹介します。

データ型の宣言

これを利用して、列の定義を指定したクラスを宣言するには、以下のように TypedDataFrame を継承したクラスを定義します。

>>> class InputData(TypedDataFrame):
...     INT_COLUMN = "int_column"
...     FLOAT_COLUMN = "float_column"
...     _dtype: Dict[str, Type] = {
...         INT_COLUMN: int,
...         FLOAT_COLUMN: float,
...     }

CSVファイルからのインスタンス生成

型定義に合わせたCSVデータ読み込みが行えるようになっています。

>>> input_data = InputData.read_csv("input_data.csv")

>>> input_data
   int_column float_column
0           1         10.0
1           2         20.0

>>> input_data.dtypes
int_column        int64
float_column    float64
dtype: object

表: input_data.csvの内容

int_column float_column
1.0 10.0
2 20.0

CSVの int_column には実数値も入っていますが、定義にしたがって整数型として読み込まれています。

コンストラクタによるインスタンス生成

コンストラクタでも型定義に合わせたデータ読み込みが行えます。

>>> InputData({"int_column": [1.0, 2.0], "float_column": [10.0, 20.0]})
   int_column float_column
0           1         10.0
1           2         20.0

データバリデーション(列の型チェック)

型定義に合わせてデータのバリデーション(列の型チェック)が行えます。

>>> InputData({"int_column": ["not_int", 2.0], "float_column": [10.0, 20.0]})
Traceback (most recent call last):
  ...
TypeError: Cannot cast array data from dtype('O') to dtype('int64') according to the rule 'safe'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
...
ValueError: invalid literal for int() with base 10: 'not_int'

int_column"not_int" という整数ではない値が入っており、列の型変換で弾かれています。

データバリデーション(列の過不足チェック)

型定義に合わせてデータのバリデーション(列の過不足チェック)が行えます。

>>> InputData({"int_column": [1.0, 2.0]})
Traceback (most recent call last):
  ...
AssertionError: 入力データの列がテーブル定義と違いますよ
入力データ: ['int_column']
テーブル定義: ['float_column', 'int_column']

渡されたデータにはint_column しか存在しないため、列の過不足チェックで弾かれています。

型ヒントとしての利用

データ処理関数の型ヒントに使えます。mypyによる静的型チェックはうまく働かせられておらず、DataFrame を返してもmypyに怒られないのですが、実行時に返り値のデータが正しい列や型を持つことのチェック(動的チェック)は行われます。

>>> class ProcessedData(TypedDataFrame):
...     ...
...
>>> class OutputData(TypedDataFrame):
...     RESULT_COLUMN1 = "result_column1"
...     RESULT_COLUMN2 = "result_column2"
...
>>> def data_processing1(input_data: InputData) -> ProcessedData:
...     ...
...     return ProcessedData(...)
...
>>> def data_processing2(processed_data: ProcessedData) -> OutputData:
...     ...
...     return OutputData(...)
...
>>> output_data: OutputData = (
...     input_data
...     .pipe(data_processing1)
...     .pipe(data_processing2)
... )

クラスメソッドによるデータ生成

作成したクラスに、そのインスタンスを生成するクラスメソッドを定義すると言ったこともできます。これによりあるクラスと別のクラスの関係を明示でき、データ処理関数を定義するよりも高い可読性を実現することができます。

>>> class OutputData(TypedDataFrame):
...     @classmethod
...     def calc(cls, input_data: InputData) -> "OutputData":
...         ...
...         return OutputData(...)
...
>>> output_data: OutputData = OutputData.calc(input_data)

まとめ

Musubi AI在庫管理のデータサイエンスロジック開発で活用している pandasの DataFrame クラスの拡張をご紹介しました。 現在はpanderaのようなオープンソースソフトウェアで概ね実現できる機能ではありますが、コンパクトな実装で様々な有用性のある機能を実現できるものではあるので紹介させていただきました。 なにかの参考になれば幸いです。

Musubi AI在庫管理のデータサイエンスロジック開発チームでは、今回ご紹介したようなちょっとエンジニアリング的な要素も、こちらの記事で議論されているようなデータサイエンス的な要素も同じチームで担当しており、エンジニアリングに強みを持つメンバーとデータサイエンス・機械学習に強みを持つメンバーが協働しながら開発を行っています。新機能のアーキテクチャ設計や新規ロジックの構築などの難易度の高い課題はそれぞれの強みを持つメンバーが行う事が多いですが、実装やロジックの精度検証のようなサブタスクに分割されたあとは、以下の様な観点を考慮してチーム内で分担しあいながら開発を行っています。

  • 各自が磨いていきたいスキルに関連するタスクが担当できるようにする
  • チーム内で知識やスキルが共有されるようにする

開発そのもの以外にも、開発プロセスや開発規約整備、スクラム開発におけるチーム運営などについて豊富な知識、経験を持ったエンジニアメンバーと議論することができる点も楽しいです。議論だけではなくエンジニアの行動から根本の考え方がにじみ出ていたりするのでとても勉強になります。

毎日チームから様々なことを学ばせてもらっていて、とても充実した日々を送っています。ご興味のある方は是非是非お問合せください。