pandas 실행시간 최적화하기

pandas 라이브러리를 사용하여 데이터를 분석할 때, 특정 코드의 실행시간이 과도하게 오래 걸리는 경우가 종종 있습니다. 그 사례와 원인은 여러 가지가 있을 텐데, 이 글에서는 제가 겪은 두 가지 상황에 대해서 소개하겠습니다. 먼저 실행 시간이 예상치 않게 오래 걸리는 경우에 대해서 살펴보고, 그 원인과 최적화 코드를 제시하는 순서로 작성하였습니다. 이 글에서 pd는 pandas 라이브러리를 호출할 때 흔히 사용되는 별칭을 의미합니다. 

 

데이터프레임에 새로운 행 추가하기

pd.concat 메서드 사용하기

첫 번째는 데이터프레임에 새로운 행(데이터)을 추가할 때 겪은 상황입니다. 데이터프레임에 10000개의 데이터를 추가하는 상황을 생각해보겠습니다. 이 때, 추가하는 데이터는 매번 새롭게 생성된다고 가정하겠습니다. 여러 개의 데이터프레임을 연결할 때는 pd.concat 메서드를 사용하는데, 새로운 데이터를 추가할 때도 이 방식을 사용할 수 있습니다. 이를 코드로 나타내면 다음과 같습니다.

 

df = pd.DataFrame(columns=["A", "B"])

for i in range(10000):
    data = pd.DataFrame({"A": [i], "B": [i * 2]})
    df = pd.concat([df, data])

 

문제는 이런 방식으로 작성한 코드는 실행시간이 과도하게 길다는 것입니다. 예시 코드는 데이터 생성 자체에 걸리는 시간이 매우 짧기 때문에 체감하기 어렵지만, 만약 데이터 생성 자체도 많은 연산을 요구한다면 코드 실행시간이 매우 길어질 것입니다. 이러한 방식은 concatenate 연산을 위한 임시 데이터프레임을 생성할 때마다 메모리를 재할당하기 때문에 메모리 측면에서도 효율적이지 않고 실행시간도 길어집니다. 이 문제는 다음과 같은 방법을 사용함으로써 해결할 수 있습니다.

 

파이썬 list 또는 dictionary 사용하기

이 방식은 생성된 데이터를 모두 모아두었다가, 데이터를 모아놓은 시퀀스 객체를 판다스 데이터프레임으로 변환하는 방법입니다. 코드로 나타내면 다음과 같습니다. 파이썬의 list 또는 dictionary를 데이터프레임으로 변환하는 구체적인 방법에 대해서는 판다스 공식문서를 참고해주세요.

 

data_list = []

for i in range(10000):
    data_list.append({"A": i, "B": i * 2})

df = pd.DataFrame(data_list)

 

실제로 두 코드를 실행하는 시간이 얼마나 차이나는지 직접 확인해보겠습니다. 앞으로 코드 실행 시간을 계속해서 확인할 것이기 때문에, 코드의 실행시간을 측정하기 위한 context manager를 하나 만들어두겠습니다. contextmanager에 대한 개념은 추후 다른 글에서 집중적으로 다뤄보겠습니다. 

 

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()
    yield
    end = time.time()
    elapsed = end - start
    print(f"Elapsed time: {elapsed} seconds")

 

이제 위에서 소개한 데이터프레임에 새로운 데이터를 추가하는 두 가지 방법에 대한 코드의 실행 시간을 측정하고 비교해보겠습니다.

 

with timer():
    df = pd.DataFrame(columns=["A", "B"])

    for i in range(10000):
        data = pd.DataFrame({"A": [i], "B": [i * 2]})
        df = pd.concat([df, data])

# Elapsed time: 3.9233627319335938 seconds


with timer():
    data_list = []

    for i in range(10000):
        data_list.append({"A": i, "B": i*2})

    df = pd.DataFrame(data_list)

# Elapsed time: 0.008012771606445312 seconds

 

pd.concat 메서드를 사용한 방법은 만 개의 데이터를 추가하는 데 약 4초가 걸렸고, 리스트를 사용하여 데이터를 한 번에 추가하는데는 0.01초 정도밖에 걸리지 않았습니다. 이 차이는 데이터의 개수가 많아질수록, 그리고 한 개의 데이터를 생성하는 시간이 길어질수록 더욱 커질 것입니다. 따라서 새로운 데이터를 추가할 때는 pd.concat 메서드 대신 데이터를 한 곳에 모아 두었다가, 한 번에 데이터프레임을 연결하는 방법을 사용해야 합니다.

 

 

pd.apply 메서드의 실행 시간 문제

다음은 pd.apply 메서드를 사용할 때 실행 시간이 과도하게 걸리던 문제에 대해서 다뤄보겠습니다. pd.apply는 정해진 축을 따라 특정한 함수를 적용하는 메서드입니다. pd.apply 메서드에 대한 자세한 내용은 공식 문서를 참조해주세요. 여기서는 최근에 실제로 겪은 상황을 예시로 pd.apply 메서드의 실행 시간 문제와 최적화 코드를 설명하겠습니다. 코드에서 사용한 데이터와 출력 결과는 저작권 문제로 표시하지 않았습니다. 

 

데이터의 컬럼 중 discount_type이라는 컬럼이 있고, 이 컬럼에는 어떤 할인을 적용받았는지에 대한 값이 담겨있습니다. 문제는 할인에 대한 정보가 숫자가 아닌 문자열이었기 때문에, 숫자만 추출해서 원래 가격을 복원해야하는 상황입니다. 예를 들어서 어떤 데이터의 discount_type 값이 "여름맞이 30% 할인 쿠폰"일 경우, 할인율인 30만을 추출하고자 합니다. 이 작업을 수행하기 위해서 먼저 다음과 같이 정규 표현식을 사용한 코드를 작성하였습니다.

 

def extract_discount_rate(text: str) -> int:
    discount_rate_pattern = r"(\d+)%"                           
    match = re.search(discount_rate_pattern, text)

    if match:
        return int(match.group(1))
    return text

 

이 함수를 pd.apply 메서드와 함께 사용하여 할인율을 추출하고, discount_rate이라는 별도의 컬럼에 저장하는 작업을 다음과 같이 수행하였습니다.

 

raw_data["discount_rate"] = raw_data.apply(
    lambda data: extract_discount_rate(data["discount_type"]), axis=1)

 

그런데 생각보다 이상할 만큼 코드 실행시간이 오래걸렸습니다. 원인은 다음과 같았습니다. 위와 같은 방식으로 apply 메서드를 적용할 경우 열 단위로 연산이 수행됩니다. 즉 extract_discount_rate 함수가 row-wise로 전체 데이터프레임에 적용됩니다. 각 행에 대하여 함수를 적용할 때마다 새로운 pd.Series 객체가 생성됩니다. 따라서 메모리 측면에서도 효율적이지 않고 실행시간도 길어지게 됩니다. 이 문제는 다음과 같이 코드를 작성하여 해결할 수 있습니다.

 

raw_data["discount_rate"] = raw_data["discount_type"].apply(extract_discount_rate)

 

이 방식은 행마다 새로운 시리즈를 생성하는 대신, 처음부터 시리즈를 선택한 다음에 pd.apply 메서드를 적용합니다. 이는 별도의 pd.Series(intermediate series)를 생성하지 않기 때문에 메모리 측면에서 효율적이며 속도도 훨씬 빠릅니다. 조금 더 기술적으로는 벡터화 개념과 관련이 있습니다.

 

Pandas 라이브러리는 행렬 연산을 최적화한 NumPy 라이브러리에 기반합니다. 넘파이는 파이썬의 반복문과 다르게 연산을 벡터화하여 행렬에 대한 연산을 한 번에 수행합니다. 첫 번째 코드는 판다스의 메서드를 사용하기는 하지만, 벡터 연산을 사용하지는 않습니다. 반면 두 번째 코드는 본체가 넘파이 배열인 판다스 시리즈 객체에 대한 벡터 연산을 사용하므로 속도가 매우 빠릅니다.

 

앞에서 정의한 context manager를 사용하여 실제로 두 코드의 실행시간이 어떠한지 비교해보겠습니다.

 

with timer():
    raw_data["discount_rate"] = raw_data.apply(
        lambda data: extract_discount_rate(data["discount_type"]), axis=1)
    
# Elapsed time: 13.987017154693604 seconds


with timer():
    raw_data["discount_rate"] = raw_data["discount_type"].apply(extract_discount_rate)

# Elapsed time: 1.3711633682250977 seconds

 

첫 번째 코드는 실행 시간이 약 14초인데 반해, 두 번째 코드는 약 1초만에 실행되는 것을 확인할 수 있습니다.

 

이처럼 같은 동작을 하는 코드라도 어떤 방식으로 작성하느냐에 따라서 실행 시간을 크게 줄일 수 있습니다. 특히 앞서 소개한 두 사례는 판다스에서 흔히 사용하는 메서드인 만큼, 최적화된 코드를 사용하면 전체 프로그램의 개발과 실행 시간을 유의미하게 단축시킬 수 있습니다.