CSV 파일과 직렬화 객체 쓰기

출력의 filesystem().open()은 쓰기 스트림을 제공합니다. 바이너리는 wb, 텍스트는 w와 인코딩을 사용하며 임의 위치에 seek하여 수정하는 방식은 지원하지 않습니다. 파일을 닫아 쓰기를 마무리합니다.

다음 세 구현은 같은 가상의 학생 자료를 csv1.csv와 csv2.csv로 나누어 저장합니다. student_id는 앞의 0을 보존하는 문자열입니다. 이 코드는 CSV 파일 자체를 저장하며 write_table()처럼 표 schema를 등록하는 경로와 구분합니다.

Polars

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


@transform.using(output=Output("/Example Project/Outputs/polars_csv_files"))
def compute(output):
    batches = {
        "csv1.csv": {
            "student_id": ["001", "002", "003"],
            "hair_color": ["brown", "blonde", "black"],
            "eye_color": ["blue", "green", "brown"],
        },
        "csv2.csv": {
            "student_id": ["004", "005", "006"],
            "hair_color": ["red", "green", "blue"],
            "eye_color": ["purple", "grey", "black"],
        },
    }
    filesystem = output.filesystem()
    for name, columns in batches.items():
        with filesystem.open(name, "wb") as stream:
            pl.DataFrame(columns).write_csv(stream)

Pandas

Pandas의 index를 별도 CSV 열로 쓰지 않도록 index=False를 지정합니다.

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


@transform.using(output=Output("/Example Project/Outputs/pandas_csv_files"))
def compute(output):
    batches = {
        "csv1.csv": {
            "student_id": ["001", "002", "003"],
            "hair_color": ["brown", "blonde", "black"],
            "eye_color": ["blue", "green", "brown"],
        },
        "csv2.csv": {
            "student_id": ["004", "005", "006"],
            "hair_color": ["red", "green", "blue"],
            "eye_color": ["purple", "grey", "black"],
        },
    }
    filesystem = output.filesystem()
    for name, columns in batches.items():
        with filesystem.open(name, "w", encoding="utf-8", newline="") as stream:
            pd.DataFrame(columns).to_csv(stream, index=False)

DuckDB

DuckDB는 로컬 임시 파일에 COPY한 다음 그 바이트를 출력 스트림으로 복사합니다. 임시 디렉터리는 출력 업로드가 끝날 때까지 유지하고 이후 정리합니다. 목적지는 이름 있는 SQL 매개변수로 전달합니다.

from pathlib import Path
import shutil
import tempfile

from transforms.api import Output, transform

QUERIES = {
    "csv1.csv": """
        SELECT * FROM (VALUES
            ('001', 'brown', 'blue'),
            ('002', 'blonde', 'green'),
            ('003', 'black', 'brown')
        ) AS t(student_id, hair_color, eye_color)
    """,
    "csv2.csv": """
        SELECT * FROM (VALUES
            ('004', 'red', 'purple'),
            ('005', 'green', 'grey'),
            ('006', 'blue', 'black')
        ) AS t(student_id, hair_color, eye_color)
    """,
}


@transform.using(output=Output("/Example Project/Outputs/duckdb_csv_files"))
def compute(ctx, output):
    connection = ctx.duckdb().conn
    filesystem = output.filesystem()
    with tempfile.TemporaryDirectory() as directory:
        for name, query in QUERIES.items():
            path = Path(directory) / name
            connection.execute(
                f"COPY ({query}) TO $destination (FORMAT CSV, HEADER true)",
                {"destination": str(path)},
            )
            with path.open("rb") as source:
                with filesystem.open(name, "wb") as destination:
                    shutil.copyfileobj(source, destination)

여기서 쿼리 본문은 모듈에 정의한 고정 SQL이며 외부 입력 문자열을 SQL에 삽입하지 않습니다. 두 파일은 모두 student_id,hair_color,eye_color 헤더와 세 데이터 행을 갖습니다.

학습한 모델 객체 저장하기

Python 객체는 pickle로 바이너리 파일에 직렬화할 수 있습니다. 다음 함수는 변환에 주입된 출력 핸들과 이미 학습한 실제 모델 객체를 인자로 받습니다.

import pickle


def write_model(output, model):
    with output.filesystem().open("model.pickle", "wb") as stream:
        pickle.dump(model, stream, protocol=pickle.HIGHEST_PROTOCOL)

복원하는 환경의 Python과 모델 라이브러리는 직렬화한 객체와 호환되어야 합니다. 신뢰하지 않는 pickle은 복원 시 코드를 실행할 수 있으므로 출처와 접근 권한을 확인한 객체만 역직렬화합니다.