[Python] 파이썬으로 특정 컬럼 기준으로 CSV 파일 분리

 

데이터 처리를 하려는 데 아래와 같은 데이터를 pd.read_csv로 읽어오려니 에러가 났다.

 

ParserError: Error tokenizing data. C error: Expected 9 fields in line 2, saw 24

 

이 데이터는 여러 헤더를 가진 데이터가 하나로 제공되고 첫번째 열(11, 15)로 데이터를 구분해야한다.

pandas는 헤더를 추론하기 때문에 계속 에러가 났고, 파이썬으로 하나의 CSV를 아예 분리하는 방법으로 생각했다.

 

코드는 Claude를 사용해 샘플데이터와 원하는 작업을 프롬프트에 작성해 얻은 것이다.

 

import csv
import os

def split_csv(input_file, output_dir, column_index):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    file_handlers = {}

    with open(input_file, 'r', newline='', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)
        
        for row in reader:
            if not row or len(row) <= column_index:  # 빈 행이거나 열이 충분하지 않으면 무시
                continue
            
            # 지정된 열의 값을 record_type으로 사용
            record_type = row[column_index].strip()
            
            # record_type이 유효한 숫자가 아니면 건너뛰기
            if not record_type.isdigit():
                continue

            # 해당 record_type의 파일 핸들러가 없으면 새로 생성
            if record_type not in file_handlers:
                output_file = os.path.join(output_dir, f'Record_{record_type}.csv')
                file_handlers[record_type] = open(output_file, 'w', newline='', encoding='utf-8')
                writer = csv.writer(file_handlers[record_type])
          
            
            # 데이터 쓰기
            writer = csv.writer(file_handlers[record_type])
            writer.writerow(row) # 해당 record_type이 들어있는 행 한줄이 저장됨. for문을 통해 모든 행을 반복하면서 진행됨

    # 모든 파일 핸들러 닫기
    for handler in file_handlers.values():
        handler.close()

# 스크립트 실행
input_file = 'sample_data.csv'
output_dir = 'split_output'
column_index = 0  # 0번째 열 (인덱스는 0부터 시작)
# sample_data는 0번째 열의 값을 기준으로 나눔. 

split_csv(input_file, output_dir, column_index)
print("CSV 파일 분리가 완료되었습니다.")

 

코드 실행시 0번째열에 고유한 값인 10, 11, 15, 21, 23 등으로 나누어서 저장됨