Daft를 활용한 고성능 데이터 파이프라인 구축 가이드: ML을 위한 Python 네이티브 데이터 엔진

import daft
from daft import col

URL = "https://github.com/Eventual-Inc/mnist-json/raw/master/mnist_handwritten_test.json.gz"
df = daft.read_json(URL)
print(df.schema())
df.show(5)

def to_28x28(pixels):
    arr = np.array(pixels, dtype=np.float32)
    if arr.size != 784:
        return None
    return arr.reshape(28, 28)

df2 = (
    df
    .with_column(
        "img_28x28", 
        col("image").apply(to_28x28, return_dtype=daft.DataType.python())
    )
)

@daft.udf(return_dtype=daft.DataType.list(daft.DataType.float32()), batch_size=512)
def featurize(images_28x28):
    out = []
    for img in images_28x28.to_pylist():
        // ...(중략)
        vec = np.concatenate([row_sums, col_sums, np.array([cy, cx, img.mean()/255.0, img.std()/255.0], dtype=np.float32)])
        out.append(vec.astype(np.float32).tolist())
    return out

df3 = df2.with_column("features", featurize(col("img_28x28")))

label_stats = (
    df3.groupby("label")
    .agg(
        col("label").count().alias("n"),
        col("pixel_mean").mean().alias("mean_pixel_mean"),
    )
    .sort("label")
)
df4 = df3.join(label_stats, on="label", how="left")

out_df = df4.select("label", "features", "pixel_mean", "pixel_std", "n")
out_path = "/content/daft_mnist_features.parquet"
out_df.write_parquet(out_path)