Pandas風ライブラリのパフォーマンス比較

この記事には[Part 2](https://qiita.com/tomohiro_ubukata/items/f702c4b505639ef9cff4)があります。
諸事情によりこちらの記事も残しておきます。

# はじめに

[SMN a.i lab. Advent Calendar 2024](https://qiita.com/advent-calendar/2024/smn-ailab) 1日目です。
一般に、機械学習では学習自体よりもデータの前処理に多大な計算リソースを消費します。
前処理ではさまざまなソフトウェアが使われますが、表形式のデータの処理においてポピュラーなライブラリとしてPandasが知られています。
しかし、近年ではPandasに近いAPIを持ちつつ、高度に最適化されたライブラリが多く配布されています。
この記事ではそれらのPandas風ライブラリのパフォーマンスを比較します。

# Pandas風ライブラリたち

## Pandas

もっともポピュラーな表形式データ処理ライブラリです。

## PyArrow

Apache ArrowのPython用ライブラリです。
それほどPandas風ではありませんが、表形式のデータ型があります。
Arrow自体は列指向のデータフォーマット仕様であり、公式から多くの言語向けにライブラリが配布されています。
PandasからPyArrowの一部の機能を利用することもできます。

## PySpark

Apache SparkのPython用ライブラリです。
それほどPandas風ではありませんが、表形式のデータ型があります。
Spark自体は主にマルチノードで分散処理することを想定したソフトウェアですが、シングルノードでも使えます。

## Polars

Pandasの代替として近年注目されているようです。
Pandasとある程度。

## Dask

Pandasとの互換性の高さを謳っています。
マルチノードでの分散処理にも対応しており、Sparkより高速だと主張しています。

## cuDF

NVIDIA GPU用のライブラリです。
"cuDF pandas Accelerator Mode" を利用するとPandasと完全な互換性があると主張しています。

# Notebook

https://colab.research.google.com/drive/1UEYHX-P9zdnwKGXH7XUK6_gpV5jph207

# 環境

Google ColaboratoryでランタイムのタイプをT4 GPUに設定して使用します。

CPUの情報

```
!lscpu
```

```
Architecture:             x86_64
  CPU op-mode(s):         32-bit, 64-bit
  Address sizes:          46 bits physical, 48 bits virtual
  Byte Order:             Little Endian
CPU(s):                   2
  On-line CPU(s) list:    0,1
Vendor ID:                GenuineIntel
  Model name:             Intel(R) Xeon(R) CPU @ 2.00GHz
    CPU family:           6
    Model:                85
    Thread(s) per core:   2
    Core(s) per socket:   1
    Socket(s):            1
    Stepping:             3
    BogoMIPS:             4000.28
    Flags:                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 cl
                          flush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc re
                          p_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3
                           fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand
                           hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp 
                          fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f a
                          vx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveop
                          t xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Virtualization features:  
  Hypervisor vendor:      KVM
  Virtualization type:    full
Caches (sum of all):      
  L1d:                    32 KiB (1 instance)
  L1i:                    32 KiB (1 instance)
  L2:                     1 MiB (1 instance)
  L3:                     38.5 MiB (1 instance)
NUMA:                     
  NUMA node(s):           1
  NUMA node0 CPU(s):      0,1
Vulnerabilities:          
  Gather data sampling:   Not affected
  Itlb multihit:          Not affected
  L1tf:                   Mitigation; PTE Inversion
  Mds:                    Vulnerable; SMT Host state unknown
  Meltdown:               Vulnerable
  Mmio stale data:        Vulnerable
  Reg file data sampling: Not affected
  Retbleed:               Vulnerable
  Spec rstack overflow:   Not affected
  Spec store bypass:      Vulnerable
  Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swa
                          pgs barriers
  Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BH
                          I: Vulnerable (Syscall hardening enabled)
  Srbds:                  Not affected
  Tsx async abort:        Vulnerable
```

ホストメモリの情報

```
!free -h
```

```
               total        used        free      shared  buff/cache   available
Mem:            12Gi       792Mi       4.0Gi       2.0Mi       7.9Gi        11Gi
Swap:             0B          0B          0B
```

GPUの情報

```
!nvidia-smi
```

```
Thu Sep 26 10:12:26 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |
| N/A   47C    P8              10W /  70W |      0MiB / 15360MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+
```

CUDA Toolkitのバージョン

```
!nvcc --version
```

```
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Aug_15_22:02:13_PDT_2023
Cuda compilation tools, release 12.2, V12.2.140
Build cuda_12.2.r12.2/compiler.33191640_0
```

# ライブラリのインストールとimport

```
!pip install pyarrow pyspark polars dask
!pip install --extra-index-url=https://pypi.nvidia.com cudf-cu12
```

```
import pandas as pd
import numpy as np

import pyarrow as pa
import pyarrow.compute as pc

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

import polars as pl

import dask.dataframe as dd

import cudf
```

# データ準備

100列 * 50万行の表形式のデータをランダムに32ビット浮動小数点数で埋め、カテゴリ列A, B, C, Dを追加します。
あまり大きすぎるとMemoryErrorになるので気を付けてください。

```
np.random.seed(42)
n_rows = 500000
n_cols = 100

# 100列 * 50万行のデータをランダムに作成
# NVIDIA GPUはたぶん32ビット浮動小数点数の方が性能がいいため32ビットで作成
data = np.random.randn(n_rows, n_cols).astype(np.float32)
df = pd.DataFrame(data, columns=[f'col_{i}' for i in range(n_cols)])

# ある列を基準にグループ化するために、適当なカテゴリ列を追加
df['group'] = np.random.choice(['A', 'B', 'C', 'D'], size=n_rows)
```

# 問題

今回はgroup列でグループ化し平均を取るという処理で比較しました。
Jupyterのtimeitマジックコマンドで10ループ * 7回実行して実行時間の平均を比較します。
また、今回はコピーの速度は主眼ではありませんが極端に差がついたのでついでに測りました。

# Pandas

Pandas

```
%%time
# 時間比較用にコピー
df_copied = df.copy()
```

```
CPU times: user 479 ms, sys: 44 ms, total: 523 ms
Wall time: 1.21 s
```

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean = df.groupby('group').mean()
```

```
The slowest run took 4.02 times longer than the fastest. This could mean that an intermediate result is being cached.
349 ms ± 219 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

# PyArrow

```
%%time
arrow_table = pa.Table.from_pandas(df)
```

```
CPU times: user 1.76 s, sys: 151 ms, total: 1.91 s
Wall time: 1.05 s
```

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean_pa = arrow_table.group_by('group').aggregate(
  [
    (f'col_{i}', "mean") for i in range(n_cols)
  ]
)
```

```
129 ms ± 34.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

倍以上速くなり、標準偏差も小さくなりました。

# PySpark

```
# Sparkセッションを作成
spark = SparkSession.builder.appName("tabular_data_processing_benchmark").getOrCreate()
```

```
%%time
df_spark = spark.createDataFrame(df)
```

```
CPU times: user 4min 39s, sys: 2.64 s, total: 4min 41s
Wall time: 4min 54s
```

今回はコピーは主眼ではありませんが、PySparkはなぜかコピーが異様に遅いです。

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean_spark = df_spark.groupBy("group").mean()
```

```
73 ms ± 32.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

5倍程度速くなり、標準偏差も小さくなりました。

# Polars

```
%%time
df_pl = pl.from_pandas(df)
```

```
CPU times: user 902 ms, sys: 138 ms, total: 1.04 s
Wall time: 661 ms
```

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean_pl = df_pl.group_by('group').mean()
```

```
80.3 ms ± 17.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

PySparkと同じくらいです。

# Dask

```
%%time
df_dask = dd.from_pandas(df, npartitions=4)
```

```
CPU times: user 694 ms, sys: 47.9 ms, total: 742 ms
Wall time: 742 ms
```

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean_dask= df_dask.groupby('group').mean()
```

```
51.6 ms ± 2.12 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

速いです。今回最速でした。

# cuDF

```
%%time
df_cudf = cudf.from_pandas(df)
```

```
CPU times: user 526 ms, sys: 123 ms, total: 649 ms
Wall time: 658 ms
```

```
%%timeit -r 7 -n 10
# group列でグループ化し平均を取る
grouped_mean_cudf = df_cudf.groupby('group').mean()
```

```
54.6 ms ± 5.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

Daskと同じくらい速いです。

# 結果

|  | 平均 |標準偏差
| ---- | ---- | ---- |
| Pandas | 349 ms | 219 ms |
| PyArrow | 129 ms | 34.3 ms |
| PySpark | 73 ms | 32.7 ms |
| Polars | 80.3 ms | 17.5 ms |
| Dask | 51.6 ms | 2.12 ms |
| cuDF | 54.6 ms | 5.16 ms |

最遅は予想通りPandasでした。
最速はDaskで、cuDFをわずかに上回りました。
最速はGPUを使うcuDFだと予想していたので、この結果は意外でした。
環境や問題の種類、測定方法によって結果は変わると思うのでいろいろな条件で比較してみてください。
特に、メモリに乗り切らないデータを分散処理する場合はPySparkとDaskの差がどうなるか気になります。

# 残課題

試している間に次から次へとPandas風ライブラリが見つかりました。
他にも以下のようなPandas風ライブラリがあるため、比較対象に追加しようと思います。

- CPU
  - Fireducks
  - Vaex
  - Modin
  - Ibis
  - Daft
- GPU
  - NVTabular
  - cuStreamz