はじめに

以前の記事にて、Pandasの条件分岐を紹介したが
より保守性と拡張性を向上する方法を見つけたので記事としました。

numpy.select()

使い所

A列の値によって、B列に入れる値を分岐させたい場合に利用する。
使用感としては、elifのあるif文のイメージである。

numpy.whereの代替として使え、比較して以下の点で優れていると感じた。

可読性が高く、後から分岐追加が容易な点
条件と選択肢で関数分けが行える点
戻り値として、新しい配列を返すので元データに変更がない点

使用例

以下のコードでは、引数のraw_dfに入っている値によって
条件分岐を行って作成した新しい配列を返している。

conditionsとchoicesの定義を別々の関数で行っているが、
まとめて1つの関数にした方が、後々修正のしやすさは高いかもしれない。

NDArrayはpandasのDataFrameの列に代入すると、
自動的にSeriesとして扱われるので型変換は不要です。

また、コード内ではマジックナンバーをそのまま記載しているが
本来であれば辞書や別ファイルで管理して、引数で渡して利用すべきである。

""" numpy.select()を用いて、整形後データを生成する処理"""

import numpy as np
from numpy.typing import NDArray

from pandas import DataFrame, Series


def _define_origin_type_conditions(raw_df: DataFrame) -> list[Series]:
    """「産地区分」列を生成するための条件を定義します。"""
    return [
        raw_df["産地国"] == "日本",
        raw_df["産地国"] != "日本",
    ]


def _create_origin_type_choices(raw_df: DataFrame) -> list[Series]:
    """「産地区分」列を生成するための選択肢を作成します。"""
    return ["国産:" + raw_df["産地国"], "輸入品:" + raw_df["産地国"]]


def generate_origin_type(raw_df) -> NDArray:
    """「産地区分」列に反映するデータを生成します。"""
    conditions: list[Series] = _define_origin_type_conditions(raw_df)
    choices: list[Series] = _define_origin_type_choices(raw_df)
    return np.select(conditions, choices, default="不明")

df.assign()

使い所

複数の列に、任意の値を入れたい際に利用する。　

以下のように列ごとに個別に反映させる方法の場合だと、
都度破壊的な変更がtransformed_dfに発生している。

対して、assign()を利用した場合だと
値が反映された後の新しいDataframeを戻り値で返すので、一度の変更で済む利点がある。

""" df.assign()を用いず、個別に値を反映する処理 """

# 生データを入れる(データ行数を確定させるため最初に設定)
transformed_df["商品名"] = raw_df["商品名"]
transformed_df["価格"] = raw_df["価格"]

# 固定値を入れる
transformed_df["カテゴリ"] = "青果"
transformed_df["保存方法"] = "冷蔵"

使用例

以下のサンプルコードでは、固定値・生データ値・整形値の区分で処理を分けており、
_generate関数で列名と値の辞書を作成し、set関数にてassign()を用いて一括で反映している。

補足として整形した値の_generate関数では、辞書のvalueで別の関数を呼び出しているが
これは前述した、numpy.select()の処理を記載した関数である。

コード内のマジックナンバーについては、numpy.select()の項で述べた内容と同じです。

""" df.assign()を用いて、一括で値を反映する処理 """

from numpy.typing import NDArray

from pandas import DataFrame, Series

# 生データを反映させるための処理
def _generate_column_raw_values(raw_df: DataFrame) -> dict[str, Series]:
    """ transformed_dfに生データを設定するための辞書を生成します。"""
    return {
        "商品名": raw_df["商品名"],
        "価格": raw_df["価格"],
    }

def set_column_raw_values(transformed_df: DataFrame, raw_df: DataFrame) -> DataFrame:
    """ transformed_dfに生データの値を入れる処理を行います。"""
    column_raw_values: dict[str, Series] = _generate_column_raw_values(raw_df)
    return transformed_df.assign(**column_raw_values)


# 固定値を反映させるための処理
def _generate_column_fixed_values() -> dict[str, str]:
    """ transformed_dfに固定値を設定するための辞書を生成します。"""
    return {"カテゴリ": "青果", "保存方法": "冷蔵"}

def set_column_fixed_values(transformed_df: DataFrame) -> DataFrame:
    """ transformed_dfに固定値の値を入れる処理を行います。"""
    column_fixed_values: dict[str, str] = _generate_column_fixed_values()
    return transformed_df.assign(**column_fixed_values)


# 整形した値を反映させるための処理
def _generate_column_transformed_values(transformed_df: DataFrame) -> dict[str, NDArray]:
    """ transformed_dfに変換データの値を入れるための辞書を生成します。"""
    return {
        "産地区分": generate_origin_type(transformed_df),
    }

def set_column_transformed_values(transformed_df: DataFrame, raw_df: DataFrame) -> DataFrame:
    """ transformed_dfに変換データの値を入れる処理を行います。"""
    column_transformed_values: dict[str, NDArray] = _generate_column_transformed_values(raw_df)
    return transformed_df.assign(**column_transformed_values)

全体の処理フロー例

上記にて紹介した、numpy.select()とdf.assign()の呼び出し元の処理を記載します。
処理ごとに関数分けしているのもあり、比較的スッキリと書けます。

""" 呼び出し元の処理 """

import pandas as pd
from pandas import DataFrame


# 元データのDataframe
raw_df: DataFrame = pd.DataFrame(
    {
        "商品名": ["りんご", "オレンジ", "バナナ"],
        "価格": [100, 80, 120],
        "産地国": ["日本", "アメリカ", "フィリピン"],
    }
)

# 変換後データのDataframe(空の状態)
transformed_df: DataFrame = pd.DataFrame(
    columns=["商品名", "価格", "カテゴリ", "保存方法", "産地区分"], dtype=str
)

# 生データを入れる(データ行数を確定させるため最初に設定)
transformed_df = set_column_raw_values(transformed_df, raw_df)

# 固定値を入れる
transformed_df = set_column_fixed_values(transformed_df)

# 変換データを入れる
transformed_df = set_column_transformed_values(transformed_df, raw_df)


print(transformed_df)

"""
   商品名   価格  カテゴリ  保存方法      産地区分
0  りんご  100    青果    冷蔵    国産:日本
1  オレンジ   80    青果    冷蔵  輸入品:アメリカ
2  バナナ  120    青果    冷蔵  輸入品:フィリピン
"""

あとがき

保守性と拡張性が高いので、Pandasを用いたデータ整形・作成作業は
numpy.select()とdf.assign()の組み合わせで行うのがベストなのではと思いつつあります。

またPandasを利用する際に、毎度同じ処理の仕方をしていては学びがないと感じたのと
最近、関数型プログラミングについて学習しておりタイムリーな内容となりました。

参考

numpy.select — NumPy v2.3 Manual
- https://numpy.org/doc/stable/reference/generated/numpy.select.html
Pandas利用者はassignメソッドを使いこなそう #Python - Qiita
- https://qiita.com/akabecos/items/fb9c1e076a77d2101dc7