Python の Pandas の iterrows は、どのぐらい遅いのか?

2024/12/30に公開

Python の Pandas の iterrows は、どのぐらい遅いのか?

前置き

Python の Pandas は、大量のデータをオンメモリで処理するとはいえ、数千万行のデータを対象とすると、 iterrows はかなり遅いようです。他の方法と比べて、どのぐらい差があるのかを調べてみました。

ここでの架空のシナリオは、在庫データを価格マスターに引き当てるというものです。DataFrame が得意とする集計や分析ではなく、その前処理部分ということになるかと思います。型も float64 ではなく Decimal を使っているなど、 DataFrame の本領を発揮するような領域での性能評価ではない ことをあらかじめお断りしておきます。

iterrows にフォーカスするため、かなり簡略化した仕様で実験しています。実際の業務では日付を保持しているでしょうし、価格も一対一対応ではないでしょうが、その辺りは考慮せず、ループの負荷をかけるのに必要な項目に絞った感じです。

データの件数は、店舗数、商品数ともに 3,000 と設定し、在庫データ、価格マスタ、ともにおよそ一千万件 (900万件) となります。また、参考までに、 Polars 版とタプルリスト版の実行時間も計測しました。

結果

何はともあれ、計測結果を見てみましょう。

計測結果

一回目

タプルリスト Pandas Polars
準備:価格辞書(dict)作成 8.523 s - -
準備:在庫データ作成 12.339 s 14.085 s 20.247 s
価格引き当て (iterrows 版) - 118.062 s 15.183 s
価格引き当て (apply 版) - 30.215 s -
価格引き当て (Series 版) - 3.343 s 6.299 s
価格引き当て (index 版) - 3.510 s -
価格引き当て (タプルリスト版) 7.555 s - -

二回目

タプルリスト Pandas Polars
準備:価格辞書(dict)作成 8.482 s - -
準備:在庫データ作成 12.166 s 13.925 s 19.441 s
価格引き当て (iterrows 版) - 114.901 s 14.904 s
価格引き当て (apply 版) - 29.661 s -
価格引き当て (Series 版) - 3.625 s 6.421 s
価格引き当て (index 版) - 3.444 s -
価格引き当て (タプルリスト版) 6.894 s - -

二回とも同じような計測値なので、大きくブレることのなく安定していると考えて良さそうです。

所感

やはり Pandas の iterrows はかなり遅いようです。最速の index を利用したものに比べて、およそ 35 倍の時間がかかっています。可能なら他の方法を選択した方が良さそうです。

iterrows の代わりとして apply を使う記事をいくつか見ました。確かに 1/4 ぐらいまで短縮できるようです。

最速は index か Series を使ったものでした。両者はだいたい同等のようです。数千万件以上のデータが対象となると、このうちのどちらかの方法を選びたいところです。

参考として計測した Polars の iter_rows は予想外の健闘で、同等の機能の Pandas の iterrows の 1/8 程度で済んでいます。 Polars には任意のカラムをインデックスにする機能はないため、 Series を使ったループが最速ということになるかもしれません。今回、初めて試してみたので、詳しくは分かりませんが…。

また、最速と予想していたタプルリストが最速でなかったことも意外でした。もちろん iterrows よりも圧倒的に速いのは間違いありません。

その他

  • 当然ではありますが、 Decimal は float64 より少し負荷がかかるようです。最初のうちは DataFrame の作成部分を float64 にして試したりもしたのですが (確か半分程度の時間で済んだハズ)、金額計算の場合 Decimal は避けられないので、 float64 での計測はバッサリ切り捨てることにしました (今回の金額引き当て処理にはあまり関係がないと思いますし)。

  • 最初、金額辞書も DataFrame で作って、 index を貼って引き当てていましたが、 Polars には任意のカラムを index にする機能がなく、 Query で引き当てていたところ、まったく処理が終わらないため、切り捨てることにしました (Pandas の index 引き当てはそこまで遅くはありませんでした)。

  • Polars 版の map_rows (Pandas#apply 相当) は Decimal を返す方法を見つけられず、断念しました。

計測環境

  • M2 MacBook Pro 2022
    • メモリ:24M
    • MacOS 15.2
  • Python 3.12.7
  • Pandas 2.2.3
  • Polars 1.18.0

ソースコード

import random as rnd
import time
from contextlib import contextmanager
from decimal import ROUND_HALF_UP, Decimal
from typing import NamedTuple

import pandas as pd
import polars as pl

PriceDict = dict[tuple[int, int], Decimal]


class StockType(NamedTuple):
    """店舗の商品在庫"""

    store_id: int
    item_id: int
    amount: Decimal


class StockAndPriceType(NamedTuple):
    """店舗の商品在庫と価格"""

    store_id: int
    item_id: int
    amount: Decimal
    price: Decimal


@contextmanager
def timer(name):
    start_tm = time.time()
    print(f"[{name}] start")
    yield
    print(f"[{name}] done in {time.time() - start_tm:.3f} s")


def make_pd_df(
    cnt_store: int,
    cnt_item: int,
    columns: list[str],
    value_range: tuple[int, int],
) -> pd.DataFrame:
    """Pandas DataFrame の作成"""
    data: list[tuple[int, int, Decimal]] = [
        (
            store_id * 100,
            item_id * 10,
            Decimal.from_float(rnd.uniform(*value_range)).quantize(
                Decimal("0.01"), ROUND_HALF_UP
            ),
        )
        for store_id in range(1, cnt_store)
        for item_id in range(1, cnt_item)
    ]

    return pd.DataFrame(data, columns=columns)


def make_pl_df(
    cnt_store: int,
    cnt_item: int,
    columns: list[str],
    value_range: tuple[int, int],
) -> pl.DataFrame:
    """Polars DataFrame の作成"""
    data: list[tuple[int, int, Decimal]] = [
        (
            store_id * 100,
            item_id * 10,
            Decimal.from_float(rnd.uniform(*value_range)).quantize(
                Decimal("0.01"), ROUND_HALF_UP
            ),
        )
        for store_id in range(1, cnt_store)
        for item_id in range(1, cnt_item)
    ]

    return pl.DataFrame(data, schema=columns, orient="row")


def make_price_dict(cnt_store: int, cnt_item: int) -> PriceDict:
    """価格の辞書作成"""
    return {
        (store_id * 100, item_id * 10): Decimal.from_float(
            rnd.uniform(0, 10000)
        ).quantize(Decimal("0.01"), ROUND_HALF_UP)
        for store_id in range(1, cnt_store)
        for item_id in range(1, cnt_item)
    }


def make_stock_pd_df(cnt_store: int, cnt_item: int) -> pd.DataFrame:
    """在庫のデータフレーム作成 (Pandas)"""
    columns = ["store_id", "item_id", "amount"]
    value_range = (0, 1000)

    df = make_pd_df(cnt_store, cnt_item, columns, value_range)

    # ダミーカラムの追加
    df["note1"] = [f"note_{i}" for i in range(len(df))]
    df["note2"] = [f"note_{i}" for i in range(len(df))]
    df["note3"] = [f"note_{i}" for i in range(len(df))]
    df["note4"] = [f"note_{i}" for i in range(len(df))]

    return df


def make_stock_pl_df(cnt_store: int, cnt_item: int) -> pl.DataFrame:
    """在庫のデータフレーム作成 (Polars)"""
    columns = ["store_id", "item_id", "amount"]
    value_range = (0, 1000)

    df = make_pl_df(cnt_store, cnt_item, columns, value_range)

    # ダミーカラムの追加
    note_df = pl.DataFrame(
        {
            "note1": [f"note_{i}" for i in range(len(df))],
            "note2": [f"note_{i}" for i in range(len(df))],
            "note3": [f"note_{i}" for i in range(len(df))],
            "note4": [f"note_{i}" for i in range(len(df))],
        }
    )

    return df.with_columns(note_df)


def make_stock_list(cnt_store: int, cnt_item: int) -> list[StockType]:
    """在庫のタプルリスト作成"""
    return [
        StockType(
            store_id * 100,
            item_id * 10,
            Decimal.from_float(rnd.uniform(0, 1000)).quantize(
                Decimal("0.01"), ROUND_HALF_UP
            ),
        )
        for store_id in range(1, cnt_store)
        for item_id in range(1, cnt_item)
    ]


def concat_pd_df1(price_dict: PriceDict, stock_pd_df: pd.DataFrame) -> pd.DataFrame:
    """価格引き当て (Pandas Series 版)"""
    prices = [
        price_dict[(store_id, item_id)]
        for store_id, item_id in zip(stock_pd_df["store_id"], stock_pd_df["item_id"])
    ]
    stock_pd_df["price"] = prices

    return stock_pd_df


def concat_pl_df1(price_dict: PriceDict, stock_pl_df: pl.DataFrame) -> pl.DataFrame:
    """価格引き当て (Polars Series 版)"""
    prices = [
        price_dict[(store_id, item_id)]
        for store_id, item_id in zip(stock_pl_df["store_id"], stock_pl_df["item_id"])
    ]
    price_df = pl.DataFrame({"price": prices})

    return stock_pl_df.with_columns(price_df)


def concat_pd_df2(price_dict: PriceDict, stock_pd_df: pd.DataFrame) -> pd.DataFrame:
    """価格引き当て (Pandas iterrows 版)"""
    prices = [
        price_dict[(row["store_id"], row["item_id"])]
        for _, row in stock_pd_df.iterrows()
    ]
    stock_pd_df["price"] = prices

    return stock_pd_df


def concat_pl_df2(price_dict: PriceDict, stock_pl_df: pl.DataFrame) -> pl.DataFrame:
    """価格引き当て (Polars iter_rows 版)"""
    prices = [
        price_dict[(row["store_id"], row["item_id"])]
        for row in stock_pl_df.iter_rows(named=True)
    ]
    price_df = pl.DataFrame({"price": prices})

    return stock_pl_df.with_columns(price_df)


def concat_pd_df3(price_dict: PriceDict, stock_pd_df: pd.DataFrame) -> pd.DataFrame:
    """価格引き当て (Pandas apply 版)"""
    prices = stock_pd_df.apply(
        lambda row: price_dict[(row["store_id"], row["item_id"])], axis=1
    )
    stock_pd_df["price"] = prices

    return stock_pd_df


def concat_pl_df3(price_dict: dict, stock_pl_df: pl.DataFrame) -> pl.DataFrame:
    """価格引き当て (Polars map_rows 版)
    未完成 (Decimal を返す方法がわからないため)
    """
    price_df = stock_pl_df.map_rows(
        lambda row: price_dict.get((row[0], row[1]), Decimal(0)),
        return_dtype=pl.Decimal(precision=None, scale=2),
    )
    # price_df = stock_pl_df.map_rows(
    #     lambda row: row[0] + row[1],
    #     return_dtype=pl.Int64,
    # )

    return stock_pl_df.with_columns(price_df)


def concat_list3(
    price_dict: PriceDict, stock_list: list[StockType]
) -> list[StockAndPriceType]:
    """価格引き当て (タプルリスト内包表記版)"""
    return [
        StockAndPriceType(store_id, item_id, amount, price_dict[(store_id, item_id)])
        for store_id, item_id, amount in stock_list
    ]


def concat_pd_df4(price_dict: PriceDict, stock_pd_df: pd.DataFrame) -> pd.DataFrame:
    """価格引き当て (Pandas index 版)"""
    prices = [
        price_dict[(store_id, item_id)] for store_id, item_id in stock_pd_df.index
    ]
    stock_pd_df["price"] = prices

    return stock_pd_df


def main():
    cnt_store = 3000
    cnt_item = 3000

    with timer("価格の辞書作成"):
        price_dict = make_price_dict(cnt_store, cnt_item)

    with timer("在庫のデータフレーム作成 (Pandas)"):
        stock_pd_df = make_stock_pd_df(cnt_store, cnt_item)

    with timer("在庫のデータフレーム作成 (Polars)"):
        stock_pl_df = make_stock_pl_df(cnt_store, cnt_item)

    with timer("在庫のタプルリスト作成"):
        stock_list = make_stock_list(cnt_store, cnt_item)

    with timer("価格引き当て (Pandas Series 版)"):
        concat_pd_df1(price_dict, stock_pd_df)

    with timer("価格引き当て (Polars Series 版)"):
        concat_pl_df1(price_dict, stock_pl_df)

    with timer("価格引き当て (Pandas iterrows 版)"):
        concat_pd_df2(price_dict, stock_pd_df)

    with timer("価格引き当て (Polars iter_rows 版)"):
        concat_pl_df2(price_dict, stock_pl_df)

    with timer("価格引き当て (Pandas apply 版)"):
        concat_pd_df3(price_dict, stock_pd_df)

    # with timer("価格引き当て (Polars map_rows 版)"):
    #     concat_pl_df3(price_dict, stock_pl_df)

    with timer("価格引き当て (タプルリスト内包表記版)"):
        concat_list3(price_dict, stock_list)

    stock_pd_df.set_index(["store_id", "item_id"], inplace=True)
    with timer("価格引き当て (index 版)"):
        concat_pd_df4(price_dict, stock_pd_df)
    stock_pd_df.reset_index(inplace=True)


if __name__ == "__main__":
    main()

Discussion