ZIP 안의 CSV를 표로 읽기

ZIP reader는 seek 가능한 입력이 필요합니다. Foundry 파일 스트림을 바로 ZipFile에 넘기지 않고 로컬 임시 파일에 복사한 뒤 엽니다. 각 ZIP 안의 .csv 파일을 읽고 디렉터리 항목과 다른 확장자는 건너뜁니다.

다음 네 구현은 헤더가 id,hair,eye,sex인 UTF-8 CSV를 가정하며 네 열을 문자열로 보존합니다. 입력 경로는 ZIP 파일을 담은 가상의 데이터셋입니다. 입력이나 일치하는 CSV가 없을 때도 지정한 네 문자열 열을 가진 빈 표를 만들 수 있습니다.

공통 ZIP 열기 함수

src/myproject/datasets/zip_helpers.py에 다음 함수를 둡니다.

from contextlib import contextmanager
import shutil
import tempfile
import zipfile


@contextmanager
def open_dataset_zip(filesystem, path):
    with filesystem.open(path, "rb") as source:
        with tempfile.TemporaryFile() as temporary:
            shutil.copyfileobj(source, temporary)
            temporary.flush()
            temporary.seek(0)
            with zipfile.ZipFile(temporary) as archive:
                yield archive

임시 파일과 ZIP은 내부 CSV를 모두 읽을 때까지 열려 있고 그 뒤 닫힙니다. ZIP 항목 이름으로 로컬 경로를 만들어 압축을 해제하지 않으므로 원본의 디렉터리 경로를 실행 환경에 재현할 필요가 없습니다.

Polars

import polars as pl
from transforms.api import Input, Output, transform
from myproject.datasets.zip_helpers import open_dataset_zip

COLUMNS = ["id", "hair", "eye", "sex"]
SCHEMA = {name: pl.String for name in COLUMNS}


@transform.using(
    source=Input("/Example Project/Inputs/zipped_students"),
    output=Output("/Example Project/Outputs/zip_polars"),
)
def compute(source, output):
    filesystem = source.filesystem()
    frames = []
    for status in filesystem.ls(glob="*.zip"):
        with open_dataset_zip(filesystem, status.path) as archive:
            for member in archive.infolist():
                if member.is_dir() or not member.filename.lower().endswith(".csv"):
                    continue
                with archive.open(member) as stream:
                    frames.append(
                        pl.read_csv(stream, infer_schema=False).select(COLUMNS)
                    )
    result = pl.concat(frames, how="vertical") if frames else pl.DataFrame(schema=SCHEMA)
    output.write_table(result)

읽은 CSV의 DataFrame을 모두 메모리에 유지하므로 결과 전체가 메모리에 들어가는 크기에 사용합니다.

Pandas

import pandas as pd
from transforms.api import Input, Output, transform
from myproject.datasets.zip_helpers import open_dataset_zip

COLUMNS = ["id", "hair", "eye", "sex"]


@transform.using(
    source=Input("/Example Project/Inputs/zipped_students"),
    output=Output("/Example Project/Outputs/zip_pandas"),
)
def compute(source, output):
    filesystem = source.filesystem()
    frames = []
    for status in filesystem.ls(glob="*.zip"):
        with open_dataset_zip(filesystem, status.path) as archive:
            for member in archive.infolist():
                if member.is_dir() or not member.filename.lower().endswith(".csv"):
                    continue
                with archive.open(member) as stream:
                    frames.append(pd.read_csv(stream, dtype="string").loc[:, COLUMNS])
    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

각 CSV를 서로 다른 임시 파일로 복사한 뒤 DuckDB의 다중 파일 reader에 경로 목록을 전달합니다. CSV 본문 문자열을 파일 경로 인자로 전달하지 않습니다. 지연 쿼리가 실제로 읽을 때까지 임시 파일을 유지합니다.

from pathlib import Path
import shutil
import tempfile

from transforms.api import Input, Output, transform
from myproject.datasets.zip_helpers import open_dataset_zip


@transform.using(
    source=Input("/Example Project/Inputs/zipped_students"),
    output=Output("/Example Project/Outputs/zip_duckdb"),
)
def compute(ctx, source, output):
    filesystem = source.filesystem()
    connection = ctx.duckdb().conn
    with tempfile.TemporaryDirectory() as directory:
        paths = []
        for status in filesystem.ls(glob="*.zip"):
            with open_dataset_zip(filesystem, status.path) as archive:
                for member in archive.infolist():
                    if member.is_dir() or not member.filename.lower().endswith(".csv"):
                        continue
                    path = Path(directory) / f"{len(paths)}.csv"
                    with archive.open(member) as content:
                        with path.open("wb") as temporary:
                            shutil.copyfileobj(content, temporary)
                    paths.append(str(path))
        if paths:
            result = connection.sql(
                """
                SELECT id, hair, eye, sex
                FROM read_csv($paths, header=true, all_varchar=true)
                """,
                params={"paths": paths},
            )
        else:
            result = connection.sql("""
                SELECT CAST(NULL AS VARCHAR) AS id,
                       CAST(NULL AS VARCHAR) AS hair,
                       CAST(NULL AS VARCHAR) AS eye,
                       CAST(NULL AS VARCHAR) AS sex
                WHERE false
            """)
        output.write_table(result)

마지막 분기는 지정한 schema의 실제 빈 relation입니다. CSV 파일은 출력의 구체화가 끝난 뒤 정리합니다. 메모리뿐 아니라 전체 CSV 임시 파일과 현재 ZIP을 담을 디스크 공간도 필요합니다.

PySpark

Spark에서는 ZIP 파일 목록을 executor에 분배하고 각 ZIP 안의 CSV 행을 flatMap으로 반환합니다. ZIP 하나의 처리는 한 작업에서 이루어집니다. CSV 파서는 인용된 쉼표와 줄바꿈을 처리하며, schema를 명시하여 결과가 0행이어도 DataFrame을 만들 수 있습니다.

import csv
import io

from pyspark.sql import types as T
from transforms.api import Input, Output, transform
from myproject.datasets.zip_helpers import open_dataset_zip

COLUMNS = ["id", "hair", "eye", "sex"]
SCHEMA = T.StructType([
    T.StructField(name, T.StringType(), True) for name in COLUMNS
])


@transform.spark.using(
    source=Input("/Example Project/Inputs/zipped_students"),
    output=Output("/Example Project/Outputs/zip_spark"),
)
def compute(source, output):
    def read_archive(status):
        with open_dataset_zip(source.filesystem(), status.path) as archive:
            for member in archive.infolist():
                if member.is_dir() or not member.filename.lower().endswith(".csv"):
                    continue
                with archive.open(member) as binary:
                    with io.TextIOWrapper(binary, encoding="utf-8", newline="") as text:
                        for row in csv.DictReader(text):
                            yield tuple(row[name] for name in COLUMNS)

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

입력이 적거나 파일 크기 편차가 크면 executor의 작업 분배도 달라집니다. 분산 파일 처리에서 파일 크기 기반 분할과 driver/executor 경계를 설명합니다.