Gunakan dengan Spark

Dokumen ini adalah pengantar singkat tentang penggunaan 🤗 Datasets dengan Spark, dengan fokus khusus pada cara memuat Spark DataFrame ke dalam objek Dataset.

Dari sana, Anda memiliki akses cepat ke elemen apa pun dan dapat menggunakannya sebagai pemuat data untuk melatih model.

Muat dari Spark

Objek Dataset adalah pembungkus dari tabel Arrow, yang memungkinkan pembacaan cepat dari array dalam dataset ke tensor PyTorch, TensorFlow, dan JAX. Tabel Arrow dipetakan ke memori dari disk, yang dapat memuat dataset yang lebih besar daripada RAM yang tersedia.

Anda dapat memperoleh Dataset dari Spark DataFrame menggunakan Dataset.from_spark():

>>> from datasets import Dataset
>>> df = spark.createDataFrame(
...     data=[[1, "Elia"], [2, "Teo"], [3, "Fang"]],
...     columns=["id", "name"],
... )
>>> ds = Dataset.from_spark(df)

Worker Spark menulis dataset ke disk dalam direktori cache sebagai file Arrow, dan Dataset dimuat dari sana.

Sebagai alternatif, Anda dapat melewati materialisasi dengan menggunakan IterableDataset.from_spark(), yang mengembalikan IterableDataset:

>>> from datasets import IterableDataset
>>> df = spark.createDataFrame(
...     data=[[1, "Elia"], [2, "Teo"], [3, "Fang"]],
...     columns=["id", "name"],
... )
>>> ds = IterableDataset.from_spark(df)
>>> print(next(iter(ds)))
{"id": 1, "name": "Elia"}

Penyimpanan sementara

Saat menggunakan Dataset.from_spark(), Dataset yang dihasilkan akan di-cache; jika Anda memanggil Dataset.from_spark() beberapa kali pada DataFrame yang sama, pekerjaan Spark yang menulis dataset sebagai file Arrow di disk tidak akan dijalankan ulang.

Anda dapat mengatur lokasi cache dengan meneruskan cache_dir= ke Dataset.from_spark(). Pastikan untuk menggunakan disk yang tersedia bagi worker Anda dan mesin Anda saat ini (driver).

Dalam sesi yang berbeda, Spark DataFrame tidak memiliki semantic hash yang sama, dan akan menjalankan ulang tugas Spark serta menyimpannya di cache baru.

Jenis fitur

Jika dataset Anda terdiri dari gambar, data audio, atau array N-dimensi, Anda dapat menentukan argumen features= di Dataset.from_spark() (atau IterableDataset.from_spark()):

>>> from datasets import Dataset, Features, Image, Value
>>> data = [(0, open("image.png", "rb").read())]
>>> df = spark.createDataFrame(data, "idx: int, image: binary")
>>> # Also works if you have arrays
>>> # data = [(0, np.zeros(shape=(32, 32, 3), dtype=np.int32).tolist())]
>>> # df = spark.createDataFrame(data, "idx: int, image: array<array<array<int>>>")
>>> features = Features({"idx": Value("int64"), "image": Image()})
>>> dataset = Dataset.from_spark(df, features=features)
>>> dataset[0]
{'idx': 0, 'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=32x32>}

Anda dapat memeriksa dokumentasi Fitur untuk mengetahui semua jenis fitur yang tersedia.