`pd.DataFrame.groupby()`のキーにすると速いデータ型・遅いデータ型

pandasのgroupbyにおける型別速度調査という記事をみかけました。

pd.DataFrame.groupby()を行ったときに、キーとする列がどのデータ型だと処理が速く、どのデータ型だと処理が遅いのかを探った記事になります。

しかし、調査対象となるデータ型はint64・float64・object・categoryのみでした。

本記事では、Int64・Float64・stringデータ型を追加してこれを行うことにしました。

import benchit
import numpy as np
import pandas as pd

CAT_N = 30
rng = np.random.default_rng(0)


def makeDataFrame(row_num):
    df = pd.DataFrame({"id_int": rng.integers(0, CAT_N, row_num)}, dtype=int)
    df = df.assign(
        id_Int64=df["id_int"].astype("Int64"),
        id_float=df["id_int"].astype(float),
        id_Float64=df["id_int"].astype("Float64"),
        id_str=df["id_int"].astype(str),
        id_string=df["id_int"].astype("string"),
        id_categorical=df["id_int"].astype("category"),
        values=rng.random(row_num),
    )
    return df

def id_int(df): df.groupby("id_int")["values"].sum()
def id_Int64(df): df.groupby("id_Int64")["values"].sum()
def id_float(df): df.groupby("id_float")["values"].sum()
def id_Float64(df): df.groupby("id_Float64")["values"].sum()
def id_str(df): df.groupby("id_str")["values"].sum()
def id_string(df): df.groupby("id_string")["values"].sum()
def id_categorical(df): df.groupby("id_categorical")["values"].sum()


inputs = [makeDataFrame(n) for n in 10 ** np.arange(8)]

funcs = [id_int, id_Int64, id_float, id_Float64, id_str, id_string, id_categorical]
t = benchit.timings(funcs, inputs)
t.plot(figsize=(10, 8), logx=True)

結果

（当然ながら？）category型の優位は動きません。それに次ぐInt64とint64、Float64とfloat64はあまり変わらないようです（なおdatetime64・timedelta64型もやってみましたがInt64とほとんど同じでした）。stringデータ型が（objectよりも！）最も遅い結果となりました。