DataFusion과 GPU를 사용하는 단일 노드 계산

표 형태의 Foundry 입력은 Pandas DataFrame, Arrow Table, Polars DataFrame 또는 원시 Parquet와 CSV 파일로 읽을 수 있습니다. 엔진에 맞는 형식을 선택하고, 지원되는 표 형식으로 결과를 write_table()에 전달합니다. 분산 Spark가 필요한지 여부는 계산 엔진 선택에서 먼저 판단합니다.

원시 파일과 출력 디렉터리

input.path()는 원시 입력 파일을 내려받은 로컬 경로를 제공합니다. 파일을 엔진에 직접 넘기면 중간 Python 객체로 역직렬화하는 과정을 줄일 수 있습니다.

직접 작성한 파일을 표 데이터셋으로 올릴 때는 다음 두 조건을 모두 지킵니다.

이 경로의 표 쓰기 API는 Parquet 파일만 받습니다.
파일을 output.path_for_write_table이 가리키는 디렉터리에 작성한 뒤 output.write_table(output.path_for_write_table)을 호출합니다.

메모리상의 Pandas, Polars와 Arrow 객체를 넘기는 방식은 파일 디렉터리를 직접 만드는 방식과 구분합니다.

Spark가 읽지 못하는 Parquet 시간 자료형

INT64 (TIMESTAMP(NANOS,false)) 같은 Parquet 시간 자료형은 Spark 버전과 지원 범위에 따라 org.apache.spark.sql.AnalysisException: Illegal Parquet type 오류를 일으킬 수 있습니다. 해당 파일을 Pandas가 읽을 수 있는 환경에서는 직접 읽어 처리합니다.

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


@transform.using(
    source=Input("/Example Project/Inputs/nanosecond_parquet"),
    output=Output("/Example Project/Outputs/converted_table"),
)
def compute(source, output):
    rows = pd.read_parquet(source.path())
    output.write_table(rows)

원래의 시간 정밀도가 필요한지, 출력 형식을 하위 엔진이 해석하는지까지 확인합니다. 입력 파일을 Pandas로 읽었다는 사실만으로 모든 하위 Spark 버전의 호환성이 결정되지는 않습니다.

Apache DataFusion

환경에 datafusion을 추가합니다. 다음 변환은 Parquet 입력의 name 문자열이 John으로 시작하는 행을 선택하여 Arrow Table로 출력합니다. 경로는 가상의 예시입니다.

from datafusion import SessionContext, col, lit
from datafusion.functions import starts_with
from transforms.api import Input, Output, transform


@transform.using(
    source=Input("/Example Project/Inputs/people"),
    output=Output("/Example Project/Outputs/john_names"),
)
def compute(source, output):
    engine = SessionContext()
    rows = engine.read_parquet(source.path())
    selected = rows.filter(starts_with(col("name"), lit("John")))
    output.write_table(selected.to_arrow_table())

to_arrow_table()로 구체화한 결과가 단일 노드 메모리에 들어가야 합니다.

cuDF에서 계산한 결과를 전달하기

cuDF는 GPU에서 DataFrame 연산을 수행합니다. 환경에 맞는 cudf 패키지를 준비하고, 플랫폼 인스턴스에 지원되는 NVIDIA GPU가 있어야 합니다. 프로젝트가 사용하는 리소스 큐에서도 해당 GPU를 사용할 수 있어야 합니다.

아래 예제는 CPU 4코어, 시스템 메모리 32GB와 NVIDIA T4를 요청합니다. 입력의 name 열을 GPU로 옮겨 John으로 시작하는 행을 정렬한 뒤, 출력 API에 전달할 Pandas DataFrame으로 변환합니다. GPU와 시스템 메모리 양쪽에 입력과 결과를 담을 여유가 필요합니다.

cuDF를 CI에서 가져오지 않고 실제 GPU 실행 시점에 가져오는 이 구성에서는 Python 하위 프로젝트 build.gradle에서 TLLV를 끕니다. TLLV는 모듈 수준 import를 전제로 하기 때문입니다.

transformsPython {
    tllv false
}

from transforms.api import Input, Output, transform


@transform.using(
    source=Input("/Example Project/Inputs/people"),
    output=Output("/Example Project/Outputs/gpu_john_names"),
).with_resources(cpu_cores=4, memory_gb=32, gpu_type="NVIDIA_T4")
def compute(source, output):
    import cudf

    rows = cudf.from_pandas(source.pandas()[["name"]])
    selected = rows[rows["name"].str.startswith("John")]
    output.write_table(selected.sort_values(by="name").to_pandas())

GPU 처리를 마친 뒤에 CPU 출력 형식으로 옮겨 불필요한 중간 왕복을 피합니다.

Polars lazy 쿼리의 GPU 실행

cudf-polars는 Polars lazy 쿼리의 GPU 실행 엔진입니다. conda_recipe/meta.yaml의 기존 run 요구사항에 패키지를 추가합니다.

requirements:
  run:
    - python
    - transforms {{ PYTHON_TRANSFORMS_VERSION }}
    - cudf-polars

GPU가 프로젝트의 리소스 큐에 제공되는지 확인한 뒤 .with_resources()로 요청합니다. 다음은 숫자 열 a, b와 분류 열 category가 있는 입력에 적용하는 예제입니다.

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


@transform.using(
    source=Input("/Example Project/Inputs/vectors"),
    output=Output("/Example Project/Outputs/category_statistics"),
).with_resources(gpu_type="NVIDIA_T4", memory_gb=32)
def compute(source, output):
    result = (
        source.polars(lazy=True)
        .with_columns(
            (pl.col("a") * pl.col("b")).alias("product"),
            (pl.col("a").pow(2) + pl.col("b").pow(2)).sqrt().alias("norm"),
        )
        .group_by("category")
        .agg(
            pl.col("product").sum().alias("total_product"),
            pl.col("norm").mean().alias("mean_norm"),
            pl.len().alias("n"),
        )
        .sort("total_product", descending=True)
    )
    materialized = result.collect(engine=pl.GPUEngine())
    output.write_table(materialized)

출력의 total_product는 분류별 곱의 합, mean_norm은 벡터 길이의 평균, n은 행 수입니다.

GPU 엔진은 Polars 연산의 일부를 지원합니다. 기본값에서는 지원하지 않는 쿼리가 CPU로 자동 전환됩니다. GPU에서 수행할 수 없으면 명시적으로 실패시켜야 하는 작업에서는 GPUEngine에 raise_on_fail=True를 지정합니다. result는 위에서 만든 동일한 LazyFrame입니다.

materialized = result.collect(engine=pl.GPUEngine(raise_on_fail=True))
output.write_table(materialized)

GPU 경로는 collect() 방식으로 전체 결과를 구체화합니다. sink_parquet() 같은 스트리밍 연산은 지원하지 않으므로 결과가 GPU와 시스템 메모리에 들어가야 합니다. CPU lazy 경로의 대용량 스트리밍 이점을 GPU 수집에도 그대로 적용하지 않습니다.