Sử dụng với PyArrow

Tài liệu này là phần giới thiệu nhanh về việc sử dụng datasets với PyArrow, tập trung đặc biệt vào cách xử lý datasets bằng các hàm tính toán Arrow, và cách chuyển đổi một dataset sang PyArrow hoặc từ PyArrow.

Điều này đặc biệt hữu ích vì nó cho phép các thao tác zero-copy nhanh chóng, do datasets sử dụng PyArrow bên trong.

Định dạng tập dữ liệu

Mặc định, các tập dữ liệu trả về các đối tượng Python thông thường: số nguyên, số thực, chuỗi, danh sách, v.v.

Để nhận các Bảng hoặc Mảng PyArrow thay thế, bạn có thể đặt định dạng của tập dữ liệu thành pyarrow bằng cách sử dụng 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"
  ]
]

Điều này cũng hoạt động với các đối tượng IterableDataset được lấy, ví dụ bằng cách sử dụng 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]]

Xử lý dữ liệu

Các hàm PyArrow thường nhanh hơn so với các hàm python viết tay thông thường, do đó chúng là một lựa chọn tốt để tối ưu hóa xử lý dữ liệu. Bạn có thể sử dụng các hàm tính toán Arrow để xử lý một dataset trong Dataset.map() hoặc 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]]

Chúng tôi sử dụng batched=True vì việc xử lý các lô dữ liệu trong PyArrow nhanh hơn so với xử lý từng hàng. Bạn cũng có thể sử dụng batch_size= trong map() để đặt kích thước của mỗi table.

Điều này cũng hoạt động với IterableDataset.map() và IterableDataset.filter().

Nhập hoặc Xuất từ PyArrow

Dataset là một lớp bọc quanh một bảng PyArrow, bạn có thể khởi tạo một Dataset trực tiếp từ bảng đó:

ds = Dataset(table)

Bạn có thể truy cập bảng PyArrow của một dataset bằng cách sử dụng Dataset.data, hàm này sẽ trả về một MemoryMappedTable hoặc một InMemoryTable hoặc một ConcatenationTable, tùy thuộc vào nguồn gốc của dữ liệu Arrow và các thao tác đã được áp dụng.

Các đối tượng đó bao bọc bảng PyArrow cơ bản, có thể truy cập tại Dataset.data.table. Bảng này chứa toàn bộ dữ liệu của dataset, nhưng cũng có thể có một ánh xạ chỉ mục tại Dataset._indices, ánh xạ chỉ mục các hàng của dataset sang chỉ mục các hàng của bảng PyArrow. Điều này có thể xảy ra nếu dataset đã được xáo trộn bằng Dataset.shuffle() hoặc nếu chỉ một tập con các hàng được sử dụng (ví dụ: sau khi gọi Dataset.select()).

Trong trường hợp tổng quát, bạn có thể xuất một dataset sang một bảng PyArrow bằng cách sử dụng table = ds.with_format("arrow")[:].