Sử dụng với Spark

Tài liệu này là phần giới thiệu nhanh về cách sử dụng 🤗 Datasets với Spark, với trọng tâm đặc biệt vào cách nạp một Spark DataFrame vào đối tượng Dataset.

Từ đó, bạn có thể truy cập nhanh bất kỳ phần tử nào và sử dụng nó như một trình nạp dữ liệu để huấn luyện mô hình.

Tải từ Spark

Một đối tượng Dataset là lớp bọc của một bảng Arrow, cho phép đọc nhanh từ các mảng trong dataset sang tensor PyTorch, TensorFlow và JAX. Bảng Arrow được ánh xạ bộ nhớ từ đĩa, cho phép tải các dataset lớn hơn bộ nhớ RAM khả dụng của bạn.

Bạn có thể lấy một Dataset từ một Spark DataFrame bằng cách sử dụng 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)

Các worker Spark ghi dataset vào đĩa trong một thư mục cache dưới dạng tệp Arrow, và Dataset được tải từ đó.

Ngoài ra, bạn có thể bỏ qua bước materialization bằng cách sử dụng IterableDataset.from_spark(), hàm này trả về một 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"}

Bộ nhớ đệm

Khi sử dụng Dataset.from_spark(), Dataset kết quả sẽ được lưu vào bộ nhớ đệm; nếu bạn gọi Dataset.from_spark() nhiều lần trên cùng một DataFrame, tác vụ Spark ghi dataset dưới dạng tệp Arrow trên đĩa sẽ không được chạy lại.

Bạn có thể đặt vị trí cache bằng cách truyền cache_dir= vào Dataset.from_spark(). Hãy đảm bảo sử dụng ổ đĩa có sẵn cho cả các worker của bạn và máy hiện tại (driver).

Trong một phiên khác, một Spark DataFrame không có cùng semantic hash, và nó sẽ chạy lại một tác vụ Spark và lưu trữ nó trong một bộ nhớ đệm mới.

Các loại tính năng

Nếu tập dữ liệu của bạn bao gồm hình ảnh, dữ liệu âm thanh hoặc mảng N chiều, bạn có thể chỉ định đối số features= trong Dataset.from_spark() (hoặc 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>}

Bạn có thể xem tài liệu Features để biết về tất cả các loại đặc trưng có sẵn.