Gunakan dengan PyArrow

Dokumen ini adalah pengantar singkat tentang penggunaan datasets dengan PyArrow, dengan fokus khusus pada cara memproses dataset menggunakan fungsi komputasi Arrow, serta cara mengonversi dataset ke PyArrow atau dari PyArrow.

Hal ini sangat berguna karena memungkinkan operasi zero-copy yang cepat, mengingat datasets menggunakan PyArrow di baliknya.

Format dataset

Secara default, dataset mengembalikan objek Python biasa: bilangan bulat, bilangan desimal, string, daftar, dan sebagainya.

Untuk mendapatkan Tabel atau Array PyArrow sebagai gantinya, Anda dapat mengatur format dataset menjadi pyarrow menggunakan Dataset.with_format():

>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("arrow")
>>> ds[0]       # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a"]]
col_1: [[0]]
>>> ds[:2]      # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
>>> ds["data"]  # pa.array
<pyarrow.lib.ChunkedArray object at 0x1394312a0>
[
  [
    "a",
    "b",
    "c",
    "d"
  ]
]

Hal ini juga berlaku untuk objek IterableDataset yang diperoleh, misalnya, menggunakan load_dataset(..., streaming=True):

>>> ds = ds.with_format("arrow")
>>> for table in ds.iter(batch_size=2):
...     print(table)
...     break
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]

Proses data

Fungsi PyArrow umumnya lebih cepat daripada fungsi python yang ditulis manual, sehingga merupakan pilihan yang baik untuk mengoptimalkan pemrosesan data. Anda dapat menggunakan fungsi komputasi Arrow untuk memproses dataset di Dataset.map() atau Dataset.filter():

>>> import pyarrow.compute as pc
>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("arrow")
>>> ds = ds.map(lambda t: t.append_column("col_2", pc.add(t["col_1"], 1)), batched=True)
>>> ds[:2]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
col_2: [[1,1]]
>>> ds = ds.filter(lambda t: pc.equal(t["col_0"], "b"), batched=True)
>>> ds[0]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["b"]]
col_1: [[0]]
col_2: [[1]]

Kami menggunakan batched=True karena memproses data dalam bentuk batch di PyArrow lebih cepat daripada baris per baris. Anda juga dapat menggunakan batch_size= di map() untuk mengatur ukuran setiap table.

Hal ini juga berlaku untuk IterableDataset.map() dan IterableDataset.filter().

Import atau Export dari PyArrow

Dataset adalah pembungkus dari Tabel PyArrow, Anda dapat membuat instans Dataset secara langsung dari Tabel:

ds = Dataset(table)

Anda dapat mengakses PyArrow Table dari sebuah dataset menggunakan Dataset.data, yang mengembalikan MemoryMappedTable atau InMemoryTable atau ConcatenationTable, tergantung pada asal data Arrow dan operasi yang diterapkan.

Objek-objek tersebut membungkus tabel PyArrow yang mendasarinya, yang dapat diakses di Dataset.data.table. Tabel ini berisi semua data dari dataset, tetapi mungkin juga ada pemetaan indeks di Dataset._indices yang memetakan indeks baris dataset ke indeks baris Tabel PyArrow. Hal ini dapat terjadi jika dataset telah diacak dengan Dataset.shuffle() atau jika hanya sebagian baris yang digunakan (misalnya setelah Dataset.select()).

Dalam kasus umum, Anda dapat mengekspor dataset ke PyArrow Table menggunakan table = ds.with_format("arrow")[:].