Использование с PyArrow

Этот документ представляет собой краткое введение к использованию datasets с PyArrow, с особым акцентом на то, как обрабатывать наборы данных с помощью вычислительных функций Arrow, а также как конвертировать набор данных в PyArrow или из PyArrow.

Это особенно полезно, так как позволяет выполнять быстрые операции без копирования данных, поскольку datasets использует PyArrow в качестве базового слоя.

Формат набора данных

По умолчанию наборы данных возвращают обычные объекты Python: целые числа, числа с плавающей точкой, строки, списки и т. д.

Чтобы получить таблицы или массивы PyArrow, вы можете задать формат набора данных как pyarrow, используя 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"
  ]
]

Это также работает для объектов IterableDataset, полученных, например, с помощью 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]]

Обработка данных

Функции PyArrow, как правило, работают быстрее обычных функций, написанных вручную на Python, поэтому они являются хорошим вариантом для оптимизации обработки данных. Вы можете использовать функции вычислений Arrow для обработки набора данных в Dataset.map() или 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]]

Мы используем batched=True, так как обработка данных в PyArrow быстрее при работе с батчами, чем построчно. Также можно использовать batch_size= в map(), чтобы задать размер каждой table.

Это также работает для IterableDataset.map() и IterableDataset.filter().

Импорт или экспорт из PyArrow

Dataset — это обёртка над таблицей PyArrow, вы можете создать экземпляр Dataset непосредственно из таблицы:

ds = Dataset(table)

Доступ к таблице PyArrow набора данных можно получить с помощью Dataset.data, которая возвращает MemoryMappedTable или InMemoryTable или ConcatenationTable в зависимости от источника данных Arrow и применённых операций.

Эти объекты оборачивают базовую таблицу PyArrow, доступную по адресу Dataset.data.table. Эта таблица содержит все данные набора, но также может существовать отображение индексов по адресу Dataset._indices, которое сопоставляет индексы строк набора с индексами строк таблицы PyArrow. Это может произойти, если набор был перемешан с помощью Dataset.shuffle() или если используется только подмножество строк (например, после Dataset.select()).

В общем случае вы можете экспортировать датасет в таблицу PyArrow с помощью table = ds.with_format("arrow")[:].