PyTorch 통합

PyTorch 통합 (Integration with PyTorch)

PyTorch 텐서는 각 결과 배치를 Apache Arrow로 내보내 DuckDB 쿼리 결과로부터 만들 수 있어요. 이 패턴은 Parquet 데이터로 모델을 학습하거나 서빙할 때 유용한데, DuckDB가 배치를 텐서로 변환하기 전에 입력을 필터링하고 프로젝션할 수 있기 때문이에요.

출처: 문서

본문

설치 (Installation)

pip install -U duckdb pyarrow torch

DuckDB에서 PyTorch로

이 예시는 Parquet 기반 Arrow 데이터셋을 쿼리하고, 결과를 RecordBatch 객체로 스트리밍하며, 각 배치를 피처(feature)와 라벨(label) 텐서로 변환해요.

import pathlib
import tempfile

import duckdb
import numpy as np
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq
import torch

base_path = pathlib.Path(tempfile.mkdtemp())
parquet_dir = base_path / "train"

table = pa.table(
    {
        "feature_0": [0.1, 0.3, 0.5, 0.7],
        "feature_1": [1.0, 0.0, 1.0, 0.0],
        "label": [0, 1, 1, 0],
    }
)
pq.write_to_dataset(table, str(parquet_dir))

con = duckdb.connect()
train_dataset = ds.dataset(str(parquet_dir))

reader = con.execute("""
    SELECT feature_0, feature_1, label
    FROM train_dataset
    WHERE label = 1
""").to_arrow_reader(batch_size=2)

for batch in reader:
    features = torch.tensor(
        np.column_stack(
            [
                batch.column(0).to_numpy(),
                batch.column(1).to_numpy(),
            ]
        ),
        dtype=torch.float32,
    )
    labels = torch.tensor(batch.column(2).to_numpy(), dtype=torch.int64)

    print(features.shape, labels)

torch.Size([2, 2]) tensor([1, 1])

DuckDB는 WHERE label = 1 필터와 선택된 컬럼을 데이터셋 스캔으로 푸시다운해서, 요청된 행과 컬럼만 텐서로 변환돼요. to_arrow_reader 메서드는 결과를 한 번에 모두 메모리에 물리지 않고 배치 단위로 스트리밍해요.

Arrow 객체를 직접 쿼리하는 방법에 대해 더 알아보려면 [“SQL on Apache Arrow” 가이드]({% link docs/current/guides/python/sql_on_arrow.md %})와 [“Export to Apache Arrow” 가이드]({% link docs/current/guides/python/export_arrow.md %})를 참고해요.

더 알아보기 (Learn more)

Arrow 통합에 대한 더 자세한 내용은 위에 링크된 가이드들을 참고해요.