스레드와 Spark executor로 파일 처리하기

단일 노드의 ThreadPoolExecutor는 같은 프로세스에서 여러 파일을 읽는 스레드를 실행합니다. Spark의 executor 분산과는 다릅니다. Spark 변환의 compute 함수는 driver에서 실행되므로 그 안에서 파일을 순회하는 것만으로는 여러 머신에 작업이 분산되지 않습니다.

다음 예제는 student_id,hair_color,eye_color 헤더가 있는 UTF-8 CSV를 읽고 세 열을 문자열로 보존합니다. 학생 ID의 앞자리 0도 보존합니다. 각 파서의 빈 값과 NA 문자열 처리 규칙은 실제 데이터 계약에 맞춰 설정합니다.

Polars 스레드 처리

from concurrent.futures import ThreadPoolExecutor

import polars as pl
from transforms.api import Input, Output, transform

COLUMNS = ["student_id", "hair_color", "eye_color"]
SCHEMA = {name: pl.String for name in COLUMNS}


@transform.using(
    source=Input("/Example Project/Inputs/student_csv_files"),
    output=Output("/Example Project/Outputs/threaded_polars"),
)
def compute(source, output):
    filesystem = source.filesystem()
    files = list(filesystem.ls(glob="*.csv"))

    def read_file(status):
        with filesystem.open(status.path, "rb") as stream:
            return pl.read_csv(stream, infer_schema=False).select(COLUMNS)

    with ThreadPoolExecutor(max_workers=4) as executor:
        frames = list(executor.map(read_file, files))
    result = pl.concat(frames) if frames else pl.DataFrame(schema=SCHEMA)
    output.write_table(result)

Pandas 스레드 처리

from concurrent.futures import ThreadPoolExecutor

import pandas as pd
from transforms.api import Input, Output, transform

COLUMNS = ["student_id", "hair_color", "eye_color"]


@transform.using(
    source=Input("/Example Project/Inputs/student_csv_files"),
    output=Output("/Example Project/Outputs/threaded_pandas"),
)
def compute(source, output):
    filesystem = source.filesystem()
    files = list(filesystem.ls(glob="*.csv"))

    def read_file(status):
        with filesystem.open(status.path, "r", encoding="utf-8", newline="") as stream:
            return pd.read_csv(stream, dtype="string").loc[:, COLUMNS]

    with ThreadPoolExecutor(max_workers=4) as executor:
        frames = list(executor.map(read_file, files))
    if frames:
        result = pd.concat(frames, ignore_index=True)
    else:
        result = pd.DataFrame({name: pd.Series(dtype="string") for name in COLUMNS})
    output.write_table(result)

DuckDB 스레드 처리

환경에 DuckDB, Pandas와 fsspec을 준비합니다. DuckDB가 StringIO 같은 Python 파일 객체를 읽을 때 fsspec이 필요합니다. 각 스레드는 자체 DuckDB 연결을 만들고 CSV 버퍼를 파싱합니다. 반환한 Pandas DataFrame은 별도의 결합 연결에 등록하여 UNION ALL로 합칩니다.

from concurrent.futures import ThreadPoolExecutor
import io

import duckdb
from transforms.api import Input, Output, transform

COLUMNS = ["student_id", "hair_color", "eye_color"]


@transform.using(
    source=Input("/Example Project/Inputs/student_csv_files"),
    output=Output("/Example Project/Outputs/threaded_duckdb"),
)
def compute(source, output):
    filesystem = source.filesystem()
    files = list(filesystem.ls(glob="*.csv"))

    def read_file(status):
        with filesystem.open(status.path, "r", encoding="utf-8") as stream:
            text = stream.read()
        with io.StringIO(text) as buffer:
            with duckdb.connect(database=":memory:") as connection:
                frame = connection.read_csv(
                    buffer, header=True, all_varchar=True
                ).fetchdf()
        return frame.loc[:, COLUMNS]

    with ThreadPoolExecutor(max_workers=4) as executor:
        frames = list(executor.map(read_file, files))
    with duckdb.connect(database=":memory:") as connection:
        if frames:
            names = []
            for index, frame in enumerate(frames):
                name = f"batch_{index}"
                connection.register(name, frame)
                names.append(name)
            sql = " UNION ALL ".join(f"SELECT * FROM {name}" for name in names)
            result = connection.sql(sql)
        else:
            result = connection.sql("""
                SELECT CAST(NULL AS VARCHAR) AS student_id,
                       CAST(NULL AS VARCHAR) AS hair_color,
                       CAST(NULL AS VARCHAR) AS eye_color
                WHERE false
            """)
        output.write_table(result)

동적으로 SQL에 들어가는 이름은 코드에서 만든 숫자 기반 테이블 이름뿐입니다. CSV 본문은 buffer로 읽고, SQL read_csv에 전달하는 문자열 경로와 구분합니다. 각 worker 결과를 메모리로 수집하므로 동시 파일 크기와 전체 결과 크기를 고려합니다.

PySpark 분산 처리

Spark FileSystem의 files()는 파일 내용 대신 FileStatus의 path, size, modified 열을 가진 DataFrame을 반환합니다. 파일 크기를 기준으로 분할된 이 목록을 executor에 보내면 각 작업이 맡은 파일을 열 수 있습니다.

다음 구현은 실제 파일을 여는 코드를 flatMap 안에서 실행합니다. CSV 헤더 이름으로 필드를 선택하므로 열 순서가 달라도 지정한 출력 열 순서로 반환합니다.

import csv

from pyspark.sql import types as T
from transforms.api import Input, Output, transform

COLUMNS = ["student_id", "hair_color", "eye_color"]
SCHEMA = T.StructType([
    T.StructField(name, T.StringType(), True) for name in COLUMNS
])


@transform.spark.using(
    source=Input("/Example Project/Inputs/student_csv_files"),
    output=Output("/Example Project/Outputs/distributed_csv"),
)
def compute(source, output):
    def read_file(status):
        with source.filesystem().open(
            status.path, "r", encoding="utf-8", newline=""
        ) as stream:
            for row in csv.DictReader(stream):
                yield tuple(row[name] for name in COLUMNS)

    files = source.filesystem().files(glob="**/*.csv")
    rows = files.rdd.flatMap(read_file)
    output.write_dataframe(rows.toDF(SCHEMA))

결과가 0행일 때 schema 없이 toDF()를 호출하면 ValueError: RDD is empty가 발생할 수 있습니다. 파일 처리에서는 위처럼 schema를 명시합니다. CSV parser가 인용된 쉼표와 줄바꿈을 해석하므로 문자열을 단순히 쉼표로 나누는 방식과 구분합니다.

파일 목록의 파티션 크기

파일을 분할할 때는 파일 크기와 파일을 여는 비용을 함께 고려합니다. 최대 크기보다 큰 단일 파일은 별도 파티션에 들어갑니다.

열기 비용은 같은 시간 동안 읽을 수 있는 바이트 수로 환산한 값이며 파일 크기에 더해집니다. 두 설정 범위의 키를 혼동하지 않습니다. Code Repositories에서 해당 Spark 설정을 담은 사용자 정의 Transforms profile을 만들고 변환에 configure()로 적용합니다.

FileSystem의 packing_heuristic에는 ffd와 wfd를 지정할 수 있습니다. 생략하면 자동 선택하며, 파일 수가 매우 많을 때 wfd는 계산이 빠른 대신 분할 균형이 덜 고를 수 있습니다. 정확한 인자와 숨김 파일 조건은 FileSystem API에 정리되어 있습니다.