ビッグデータをPythonとApache Arrowで扱ってみた

はじめに

業務でビッグデータを扱うことになり、データの読み込みとプログラム間の受け渡しに課題が出てきたことから、なにか良いアイディアがないかなと思っていたらApache Arrowというのがあるとご紹介頂いたので試してみた。

試してみたこと

機械学習のためにデータを読み込むのですが、ある特定のCSVデータを繰り返し読み込む必要があるので、このCSVをApache Arrow形式に変換して保存し、ファイルの読み込みを高速化できるか試してみる。

pyarrow

PythonでApache Arrowを使用するには、pyarrowというライブラリを使うらしい。
これは、以下のコマンドでインストール。

pip install pyarrow

CSVデータ読込

キー、日時、値の3つの列があるデータが100,000件入っていて、サイズが42MBのファイルが10個あり、以下のコマンドでCSVファイルの読み込み時間を計測。

read_csv.py

import glob
import pyarrow as pa
from time import time

files = glob.glob('data/*.csv')

df_org = pd.DataFrame()

start_time = time()

for file in files:
    df_org = pd.concat([df_org, pd.read_csv(file)])
        
df_org = df_org.reset_index(drop=True)

print("read csv : {0}s".format(time() - start_time))

上記スクリプトを実行すると以下の結果となりました。

read csv : 13.279139041900635s

Arrowデータ書出

import glob
import pyarrow as pa
from time import time

files = glob.glob('data/*.csv')

df_org = pd.DataFrame()

for file in files:
    df_org = pd.concat([df_org, pd.read_csv(file)])
        
df_org = df_org.reset_index(drop=True)

start_time = time()

table = pa.Table.from_pandas(df_org)
record_batch = table.to_batches()

with pa.OSFile("pw.arrow","wb") as sink:
    
    schema = record_batch[0].schema
    writer = pa.RecordBatchFileWriter(sink, schema)
    writer.write_batch(record_batch[0])
    writer.close()
    
print("write arrow : {0}s".format(time() - start_time))

上記のスクリプトを実行すると以下のとおり。

write arrow : 7.5772552490234375s

Arrowデータ読込

書き出したArrowデータを読み込んでみる

import pyarrow as pa
from time import time

start_time = time()

reader = pa.RecordBatchFileReader("pw.arrow")
rb = reader.get_record_batch(0)
df_aw = rb.to_pandas()

print("read arrow : {0}s".format(time() - start_time))

上記のスクリプトを実行すると以下のとおり。

read arrow : 3.1053459644317627s

考察

1,000,000件のデータをCSV形式から読み込んだ場合とArrow形式(?)で読み込んだ場合の比較

形式	時間[s]
CSV	13.28
Arrow	3.11

...圧倒的ですね(^_^;)

CSV形式とArrow形式でのファイルサイズの比較

形式	サイズ[MB]
CSV	420
Arrow	600

...約1.4倍ですね(^_^;)

この差をどう捉えるかですが、新しく入ってきたCSVデータを一回Arrow形式に変換して保存することにより、次からの読込時間は圧倒的に早くなりそうですね。

とりあえず、今の仕事では反応速度の方が重要なので、Arrow形式を使ってデータを保存してみたいと思います。

以上、現場からきむらが報告いたしました。