[python] 제너레이터의 필요성(메모리 효율성)

제너레이터란?

파이썬에는 이터레이터(iterator)를 생성할 때 사용되는 제너레이터(generator)라는 것이 있습니다. yield라는 키워드가 제너레이터와 밀접한 관련이 있습니다. 어떠한 함수 안에 yield가 있다면 해당 함수는 제너레이터 함수입니다. 제너레이터 객체는 함수의 코드를 조금씩 실행할 때 사용됩니다. 결과적으로 메모리를 적게 사용합니다. 그와 관련된 예제를 하나 살펴보시면 제너레이터의 필요성을 느끼실 수 있을 것입니다.  

 

우선 아주 많은 행, 약 1억 행으로 이뤄진 csv 파일의 행 개수를 파악하는 코드를 짜야한다고 가정해보겠습니다. 저는 아래 링크에서 해당 csv 파일을 얻었습니다. 일본의 무역 관련 데이터라고 하는데, 지금 데이터의 내용은 전혀 중요하지 않습니다. 그냥 행이 무지하게 많은 csv 파일이라는 것이 중요합니다. 

 

https://www.kaggle.com/datasets/4e614ec846ab778f6a2ff166232d5a65f5e6786b4f5781690588bd2cccd71cb6?resource=download&select=custom_1988_2020.csv 

 

제너레이터 미사용

먼저 제너레이터를 사용하지 않고, csv 파일을 한 번에 읽어서 각 행을 하나의 리스트에 담아서 리턴해주는 함수를 사용하여 해당 기능을 구현해보겠습니다. 

 

def csv_reader(file_name):
    print("csv_reader 함수 실행")
    file = open(file_name)
    result = file.read().split("\n")
    return result


csv_gen = csv_reader("custom_1988_2020.csv")
print(type(csv_gen))

row_count = 0

for row in csv_gen:
    row_count += 1

print(f"Row count is {row_count}")

 

위 코드를 실행하면 다음과 같이 함수 내부에 있는 print 문을 통해서 "csv_reader 함수 실행"이라는 문구가 터미널에 출력되고, 행 개수도 잘 계산해냅니다. 일단 우리가 원하는 기능은 수행해냈습니다. 

 

 

하지만, 작업관리자에서 메모리 사용량 그래프를 확인해보시면 메모리 사용량이 크게 증가하는 것을 알 수 있습니다. 제 PC가 32GB의 RAM을 사용하고 있어서 그나마 감당할 수 있는 것이지 저사양의 PC에서는 100%을 꽉 채워버릴 것입니다. 

 

 

제너레이터 사용

이번에는 제너레이터를 사용해서 동일한 기능을 수행하는 코드를 작성해보겠습니다. 

 

def csv_reader(file_name):
    print("csv_reader 함수 실행")
    for row in open(file_name, "r"):
        yield row


csv_gen = csv_reader("custom_1988_2020.csv")
print(type(csv_gen))
row_count = 0

for row in csv_gen:
    row_count += 1

print(f"Row count is {row_count}")

 

제너레이터를 사용해도 역시 행 개수를 파악해내는 우리의 요구사항을 잘 수행합니다. 

 

 

그리고 제너레이터를 사용할 때 놀랍게도 메모리 사용량이 평정심을 유지한다는 것을 확인할 수 있습니다. 

 

 

참고자료

[1] https://realpython.com/introduction-to-python-generators/