Integration mit PyTorch

PyTorch-Tensoren können aus DuckDB-Abfrageergebnissen erstellt werden, indem jeder Ergebnis-Batch als Apache Arrow exportiert wird. Dieses Muster ist nützlich beim Trainieren oder Bereitstellen von Modellen auf Parquet-Daten, weil DuckDB die Eingabe filtern und projizieren kann, bevor die Batches in Tensoren umgewandelt werden.

Installation

pip install -U duckdb pyarrow torch

DuckDB nach PyTorch

Dieses Beispiel fragt ein Parquet-basiertes Arrow-Dataset ab, streamt das Ergebnis als RecordBatch-Objekte und wandelt jeden Batch in Feature- und Label-Tensoren um.

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 schiebt den Filter WHERE label = 1 und die ausgewählten Spalten in den Dataset-Scan, sodass nur die angeforderten Zeilen und Spalten in Tensoren umgewandelt werden. Die Methode to_arrow_reader streamt das Ergebnis in Batches, anstatt das gesamte Ergebnis auf einmal im Speicher zu materialisieren.

Weitere Informationen zum direkten Abfragen von Arrow-Objekten finden Sie im Leitfaden „SQL auf Apache Arrow“ und im Leitfaden „Export nach Apache Arrow“.