Python. 대용량 파일을 스트림 프로세스로 분할하기

 

 

대용량 파일을 분할하려고 하는데, 통으로 읽어서 하려다 자꾸 OOM이 나서 스트림 처리 분할 로직을 만들어 봤습니다.

매번 자바만 하다가 파이썬을 하니 굉장히 재밌더라고요. 대학생 때로 돌아간 것 같고..

 

# 정책

스택 드라이버 코드를 참고해서 만들었고, 기본 동작 방식은 다음과 같습니다.

 

1. `block_size` 크기만큼 파일을 읽습니다. --> `block`

2. `block`에서 `marker`가 나올 때마다 해당 값들을 `rows`에 넣습니다.

3. `rows`가 특정 크기 이상되면 파일을 씁니다.

 

# 코드

def split_file():
    marker = '{delimiter}'
    partition_no = 0
    max_row_cnt = 500000
    block_size = 8192
    encoding = 'utf-16'

    with open("my_file.txt", encoding=encoding, mode='r') as fr:
        rows = []
        current = ''
        for block in iter(lambda: fr.read(block_size), ''):
            current += block.replace('\x00', '').replace('\n', '\\n').replace('\r', '\\r')

            while True:
                marker_pos = current.find(marker)
                if marker_pos == -1:
                    break
                rows.append(current[:marker_pos])
                current = current[marker_pos + len(marker):]

                if len(rows) >= max_row_cnt or (len(block) < block_size and marker not in current):
                    with open(f"my_files_PART_{partition_no}.txt", encoding=encoding, mode='w') as fw:
                        fw.writelines("%s\n" % line for line in rows)
                    partition_no += 1
                    rows.clear()

        if partition_no == 0:
            print(f"## Empty data.. {'file..'}")
            with open(f"my_files_PART_{partition_no}.txt", encoding='utf-16', mode='w') as fw:
                fw.writelines(" \n")

        rows.append(current)

 

# 후기

간간히 자바 이외의 언어도 하면서 굳은 머리를 풀어주어야 하겠습니다.