[파이썬] Dataframe 대용량 빅데이터 빠르게 읽고 쓰는 방법

파이썬을 사용해서 수 GB 정도의 데이터를 다루는 경우가 많습니다. 이 정도 크기의 대용량데이터를 판다스를 사용해서 읽고 쓰면 상당히 많은 시간이 소요됩니다. 40~50초 정도에서 길면 몇십 분 정도 기다려야 하는 경우가 있죠. 조금 기다리면 되지만 저는 가능하면 빠르게 처리되는 것을 선호합니다. 그래서 새롭게 알게 된 pyarrow를 사용해서 더 빠르게 데이터 프레임 형식으로 데이터를 읽고 쓸 수 있는 방법에 대해서 알아보겠습니다.

 

Apache Arrow란?

Apache Arrow는 인메모리 분석을 위한 개발 플랫폼으로 빅데이터 저장 및 처리 등을 지원하는 기술이 포함되어 있습니다. Arror 객체는 C++를 통해서 구현되어 있고 NumPy, Pandas 등의 Python 객체가 내장되어 통합이 용이합니다.

Pyarrow는 table이라는 형식을 사용합니다. Pandas는 DataFrame을 기본 형식으로 사용하는 것과 비슷하게 생각하면 됩니다.

 

Pyarrow 설치

pip을 사용해서 쉽게 설치할 수 있습니다.

pip install pyarrow

 

Test용 임시 데이터 만들기

테스트를 위해 임시로 간단한 데이터를 만들어서 사용하겠습니다. 기존에 사용할 데이터가 있다면 해당 데이터를 사용해도 무방합니다.

import os
import pandas as pd
import random

save_dir = os.path.join(os.getcwd(), 'tmp')
os.makedirs(save_dir, exist_ok=True)
file_path = os.path.join(save_dir, '임시데이터.csv')
line_count = 100_000

with open(file_path, 'w') as handler:
    start = time.perf_counter()
    for _ in range(line_count):
        data_generated = ["ABCDEFG"[random.randint(0,6)] + 
                          str(random.randint(1000, 9999)) for _ in range(1000)]
        line_data = ",".join(data_generated) + "\n"
        handler.write(line_data)
    print(time.perf_counter()-start)

만들어서 저장한 데이터

 

데이터 형식은 문자열로 object type으로만 만들어서 구성하였습니다. 뒤에서 판다스와 아파치를 사용한 경우의 시간 차이를 비교해보겠지만, int type으로만 구성된 데이터의 경우는 pandas와 pyarrow의 차이가 크게 없는 것으로 결과가 나왔습니다.

Pyarrow CSV 파일 읽기

pyarrow 패키지 내에 csv 파일 형식을 다루는 모듈이 있습니다. csv.read_csv를 통해서 csv파일을 읽어올 수 있습니다. 큰 틀은 판다스와 비슷합니다.

from pyarrow import csv

table = csv.read_csv(file_path)
table

read_csv, write_csv를 할 때 pyarrow는 옵션을 pyarrow의 옵션을 따로 묶어서 넘겨줘야합니다. 판다스와 비교한 아래 예시를 참고해 주세요.

Pandas

import pandas as pd

df = pd.read_csv(file_path, columns=['a','b','c',])

 

Pyarrow

from pyarrow import csv

options = csv.ReadOptions(column_names=['a','b','c'])
table = csv.read_csv(file_path, read_options=options)

비교하면 쉽게 차이가 보이죠. csv.ReadOptions를 통해서 columns명을 정해서 넘겨줄 수 있습니다. 이렇듯 pyarrow에서는 옵션을 하나로 모아서 넘겨줘야 한다는 차이가 있습니다.

 

Pyarrow CSV 파일 쓰기

from pyarrow import csv

options = csv.WriteOptions(include_header=False)
csv.write_csv(table, "data.csv", options)

 

Pandas VS Pyarrow 읽기 시간 비교

판다스와 Pyarrow를 사용해서 데이터 읽고 쓰는 시간을 비교해 보겠습니다. 최종적으로 Pandas의 DataFrame을 사용해서 데이터 분석이나 시각화를 진행한다고 생각해서 pyarrow의 table 형태를 dataframe으로 변경하는 것까지 해서 시간을 비교해 보겠습니다.

import time
import os
import pandas as pd
from pyarrow import csv

print("Pandas DataFrame 사용")
start = time.perf_counter()

df = pd.read_csv(file_path, header=None)
print(df.shape)
print(df.head())

end = time.perf_counter()
print('소요시간 = {:.2f} sec'.format(end-start), end="\n")


print("Pyarrow Table 사용")
start = time.perf_counter()

read_option = csv.ReadOptions(column_names=list(str(i) for i in range(1000)))
table = csv.read_csv(file_path, read_options=read_option)
df = table.to_pandas()
print(df.shape)
print(df.head())

end = time.perf_counter()
print('소요시간 = {:.2f} sec'.format(end-start))

 

판다스를 통해서 pd.read_csv를 사용한 경우 23.45 sec 가 소요됐습니다. 반면 Pyarrow의 csv.read_csv를 사용해서 table 형식으로 읽은 후 다시 DataFrame으로 변경하는 데 걸리는 시간은 11.58 sec로 거의 50%가량의 시간을 단축시킬 수 있었습니다.

Pandas와 Pyarrow 읽기 속도 비교 결과

 

Pandas VS Pyarrow 쓰기 시간 비교

데이터프레임 형식의 데이터를 CSV 형식으로 저장하는 데 걸리는 시간을 비교해 보겠습니다. 

import time
import os
import pandas as pd
from pyarrow import csv
import pyarrow as pa

# 쓰기 전 데이터 읽어오기
# 읽는 것은 Pyarrow가 빠르니까 pyarrow 사용
read_option = csv.ReadOptions(column_names=list(str(i) for i in range(1000)))
table = csv.read_csv(file_path, read_options=read_option)
df = table.to_pandas()


# 쓰기 시간 비교
print("Pandas DataFrame 사용")
start = time.perf_counter()

df.to_csv(file_path, header=None)

end = time.perf_counter()
print('소요시간 = {:.2f} sec'.format(end-start), end="\n")


print("Pyarrow Table 사용")
start = time.perf_counter()

options = csv.WriteOptions(include_header=False)
table = pa.Table.from_pandas(df)
csv.write_csv(table, file_path, options)

end = time.perf_counter()
print('소요시간 = {:.2f} sec'.format(end-start))

Pandas vs Pyarrow 쓰기 시간 비교 결과

쓰기 결과가 상당히 차이가 많이 나는 것을 확인했습니다. Pandas는 45초가 걸린 반면 Pyarrow는 7초밖에 걸리지 않았습니다. 예상보다 훨씬 큰 차이를 보였습니다. 거의 6배 정도의 차이니까 엄청난 차이입니다. 이 정도면 Pyarrow를 무조건 써야 하지 않겠습니까?!

 

지금까지 대용량데이터를 어떻게 하면 더 빠르게 DataFrame 형식으로 읽어오거나 CSV로 저장할 수 있는지에 대해서 알아보았습니다. 해당 방법을 활용해서 더욱 빠른 파이썬 활용에 도움이 됐으면 좋겠습니다. Pyarrow를 잘 활용해 보세요~

 

참고하기 좋은 자료

https://superfastpython.com/multithreaded-file-loading/#Create_10000_Files_to_load