(pthon) pandas 처리 속도 개선에 대하여

지금 하고 있는 프로젝트에서 python pandas 와 numpy를 사용해서 데이터를 처리하는 부분이 있다.
그런데 기존에 있던 처리 속도가 느려서 이 부분을 검토하면서,
처리 속도를 높이기 위해 이것 저것 수정해서 테스트 하다 보니 알게 된 것들이 있어서 정리해 본다.

- 최선의 방식은 numpy ndarray vectorization 처리다.

- for loop를 이용해서 모든 데이터를 순회하는 방식은 피해야 한다. 
  (nested for loop 방식은 어떤 식 으로 든 개선의 여지가 있다)

그럼 위 내용에 대해서 상세히 확인 해 보도록 하자.
우선 기존 소스 분석을 통해서 로직을 심플 하게 추려서 만들어 보면 다음과 같았다.

import pandas as pd
import numpy as np
import sys
import time

loop_result =[]
df_big   = pd.DataFrame(np.random.random_sample((5000, 2)), columns=["X", "Y"])
df_small = pd.DataFrame(np.random.random_sample((  50, 3)), columns=["x_val", "y_val", "dummy"])

start = time.time()
for index_small in df_small.index:
    list_distance = []
    for index_big in df_big.index:
        distance = (df_small.loc[index_small, "x_val"] - df_big.loc[index_big, "X"]) ** 2 + (
            df_small.loc[index_small, "y_val"] - df_big.loc[index_big, "Y"]
        ) ** 2
        list_distance.append(distance)

    if len(list_distance) > 0:
        min_index = list_distance.index(min(list_distance))
        loop_result.append(min_index)

end = time.time()
print("elapsed (for loop)      = ", end - start)

형태(shape)가 서로 다른 두 개의 dataframe을 가지고 for loop를 돌면서,
서로 간의 거리를 모두 구해서 그중에 제일 작은 값을 구하는 처리를 수행하고 있었다.

보시다시피 df_big 과 df_small 은 shape 가 서로 다르다 (5000 x 2 vs 50 x 3).
df_small 의 모든 row 마다 df_big 의 모든 row 간 거리(squared euclidean distance)를 구해서 그중 min 값들을 구하고 있다.

이것을 수행 해보면 내 pc 에서는 대략 6초 정도 걸린다.
이 정도면 별거 아니라고 생각 할 수도 있겠지만 실제 코드에서 처리되는 데이터는 
5000 x 50개가 아니고 수만 x 수천 단위이며 또 이런 식으로 처리되는 것이 
반복해서 수행되기도 해서 결국 프로젝트에 있던 소스는 처리 시간이
수 분 이상 걸리고 있었다.

개선 시도 1 (iterrows)

처음 시도 했던 것은 for loop 부분에서 iterrows()를 사용해보는 것이었다.

for _, small_row in df_small.iterrows():
        list_distance = []
        for _, big_row in df_big.iterrows():
            distance = (small_row["x_val"] - big_row["X"]) ** 2 + (small_row["y_val"] - big_row["Y"]) ** 2
            list_distance.append(distance)
        #... 상동 ...

결과는 8초 정도로, 오히려 더 느려졌다. 인터넷 검색해보면 for 보다는 좀 빠르게 나온다고 
하는데 이 경우엔 해당되지 않았다.
즉, 이걸 쓰면 무조건 빨라 진다는 건 없는 거 같다. 
절대적이지 않고 테스트 해봐야만 알 수 있다.

개선 시도 2 (at)

그럼 for loop 부분에서 loc 보다는 좀 더 빠르다는 at 으로 변경 해보자.

for index_small in df_small.index:
        list_distance = []
        for index_big in df_big.index:
            distance = (df_small.at[index_small, "x_val"] - df_big.at[index_big, "X"]) ** 2 + (
                df_small.at[index_small, "y_val"] - df_big.at[index_big, "Y"]
            ) ** 2
            #... 상동 ...

결과는 3초 정도로 많이 빨라졌다.

개선 시도 3 (vectorization)

지금 이중 for loop 를 돌고 있는데 안쪽 for loop 대신 벡터 연산 으로 대체 해 보자.

for index_small in df_small.index:
        list_distance = (df_small.at[index_small, "x_val"] - df_big["X"]) ** 2 + (
            df_small.at[index_small, "y_val"] - df_big["Y"]
        ) ** 2

        if len(list_distance) > 0:
            min_index = list_distance.argmin() 
            loop_result.append(min_index)

df_big 에 대해서 개별 row 접근이 아닌 벡터 연산 처리로 수정을 했다.
실행 해보면 이 간단한 수정으로 0.03 초로 감소되어 무려 200배 정도 빨라졌다. 

개선 시도 4 (vectorization + numpy)

0.03초 면 충분하다고 할 수도 있지만, 실제 데이터는 수십 만 건일 수도 있어서 안심할 수 없었다. 
여기서 조금 더 나아가 보자.
이번엔 df_big 를 numpy array 로 변경해서 처리 해보자.

for index_small in df_small.index:
        list_distance = (df_small.at[index_small, "x_val"] - df_big["X"].to_numpy()) ** 2 + (
            df_small.at[index_small, "y_val"] - df_big["Y"].to_numpy()
        ) ** 2
        #... 상동 ...

0.03 초에서 0.00502 초로 6배 정도 더 빨라졌다.
numpy 배열은 동일한 type으로 처리되기 때문에 불필요한 동적 type 체크 등을 
수행하지 않아, pandas series 처리보다 훨씬 빠르다.

개선 시도 5 (apply + vectorization + numpy)

코드엔 아직 for loop 가 하나 더 남아있는데, 인터넷에 검색해보면 나오는 
수많은 pandas speed tricks 중에 이런 for loop 사용은 지양 해야 할 순위 1 번으로 되어 있다.
그렇다면 이것을 apply 처리로 한번 수정해 본다.

apply_result =[]

start = time.time()
df_small.apply(get_distance, axis=1, args=[df_big])
end = time.time()
print("elapsed (apply) = ", end - start)
return end - start
    
def get_distance(small_row, arg_df_big):
    x_val = small_row["x_val"]
    y_val = small_row["y_val"]
    list_distance = (x_val - arg_df_big["X"].to_numpy()) ** 2 + (y_val - arg_df_big["Y"].to_numpy()) ** 2
    
    if len(list_distance) > 0:
        min_index = list_distance.argmin()
        apply_result.append(min_index)

수행 시 결과는 0.00273 초 가 나와서, 약간 더 빨라졌다.

위 테스트 내용을 바탕으로 기존 코드를 수정 해보니, 
예제처럼 수백 수천 배 빨라졌으면 좋았겠지만 전체 코드를 모두 수정할 수 없는 상황이어서 
그 정도의 개선을 이루진 못했다. 기존 코드 일부분만 수정을 했고, 그 정도의 작업으로도 
전체 수행 시간을 1시간 20분 --> 10분으로 단축 시킬 수 있었다.

전체 코드 

https://gist.github.com/jeremyko/b92399a94dd01db3cf5499a3a2b5afc4