資料集文件
與 Spark 搭配使用
並獲得增強的文件體驗
開始使用
與 Spark 一起使用
本文件簡要介紹如何將 🤗 Datasets 與 Spark 一起使用,特別著重於如何將 Spark DataFrame 載入至 Dataset 物件。
載入後,您可以快速存取任何元素,並將其作為資料載入器(data loader)來訓練模型。
從 Spark 載入
Dataset 物件是對 Arrow 資料表的封裝,它允許將資料集中的陣列快速讀取為 PyTorch、TensorFlow 和 JAX 張量。Arrow 資料表是從磁碟進行記憶體映射(memory mapped)的,因此可以載入大於可用 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 工作節點(workers)會將資料集以 Arrow 檔案的形式寫入磁碟的快取目錄中,並從該處載入 Dataset。
或者,您可以使用 IterableDataset.from_spark() 來跳過實體化(materialization)過程,這會回傳一個 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"}快取(Caching)
當使用 Dataset.from_spark() 時,產生的 Dataset 會被快取;如果您對同一個 DataFrame 多次呼叫 Dataset.from_spark(),它不會重新執行將資料集以 Arrow 檔案寫入磁碟的 Spark 作業。
您可以透過傳遞 cache_dir= 給 Dataset.from_spark() 來設定快取位置。請務必使用工作節點與當前機器(驅動程式)皆可存取的磁碟。
在不同的連線(session)中,Spark DataFrame 不會有相同的 語意雜湊值(semantic hash),因此它會重新執行 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 上更新