與 Spark 共同使用

本文件簡要介紹瞭如何將 🤗 Datasets 與 Spark 結合使用，特別關注如何將 Spark DataFrame 載入到 Dataset 物件中。

之後，您就可以快速訪問任何元素，並將其用作資料載入器來訓練模型。

從 Spark 載入

Dataset 物件是 Arrow 表的包裝器，它允許從資料集中的陣列快速讀取資料到 PyTorch、TensorFlow 和 JAX 的張量中。Arrow 表是從磁碟進行記憶體對映的，這可以載入比您可用 RAM 更大的資料集。

您可以使用 Dataset.from_spark() 從 Spark DataFrame 獲取一個 Dataset 物件。

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

Spark 工作節點將資料集以 Arrow 檔案的形式寫入磁碟上的快取目錄中，然後 Dataset 從那裡載入。

或者，您可以透過使用 IterableDataset.from_spark() 來跳過物化過程，該方法會返回一個 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"}

快取

當使用 Dataset.from_spark() 時，生成的 Dataset 會被快取；如果您在同一個 DataFrame 上多次呼叫 Dataset.from_spark()，它不會重新執行將資料集寫入磁碟為 Arrow 檔案的 Spark 作業。

您可以透過向 Dataset.from_spark() 傳遞 cache_dir= 來設定快取位置。請確保使用對您的工作節點和當前機器（驅動程式）都可用的磁碟。

在不同的會話中，Spark DataFrame 不具有相同的語義雜湊值，它會重新執行一個 Spark 作業並將其儲存在一個新的快取中。

特徵型別

如果您的資料集由影像、音訊資料或 N 維陣列組成，您可以在 Dataset.from_spark() (或 IterableDataset.from_spark()) 中指定 features= 引數。

>>> 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>}

您可以查閱 Features 文件來了解所有可用的特徵型別。

< > 在 GitHub 上更新