Utiliser avec PyArrow

Ce document est une introduction rapide à l’utilisation de datasets avec PyArrow, en mettant particulièrement l’accent sur la manière de traiter les jeux de données à l’aide des fonctions de calcul Arrow, ainsi que sur la conversion d’un jeu de données vers PyArrow ou depuis PyArrow.

C'est particulièrement utile car cela permet des opérations rapides sans copie, datasets utilisant PyArrow en arrière-plan.

Format du jeu de données

Par défaut, les jeux de données renvoient des objets Python classiques : des entiers, des flottants, des chaînes de caractères, des listes, etc.

Pour obtenir des Tables ou des Arrays PyArrow à la place, vous pouvez définir le format du jeu de données sur pyarrow à l'aide de 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"
  ]
]

Cela fonctionne également pour les objets IterableDataset obtenus, par exemple, à l'aide de 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]]

Traiter les données

Les fonctions PyArrow sont généralement plus rapides que les fonctions Python écrites à la main, et constituent donc une bonne option pour optimiser le traitement des données. Vous pouvez utiliser les fonctions de calcul Arrow pour traiter un jeu de données dans Dataset.map() ou 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]]

Nous utilisons batched=True car il est plus rapide de traiter des lots de données dans PyArrow que ligne par ligne. Il est également possible d’utiliser batch_size= dans map() pour définir la taille de chaque table.

Cela fonctionne également pour IterableDataset.map() et IterableDataset.filter().

Importer ou exporter depuis PyArrow

Un Dataset est un conteneur d'une table PyArrow, vous pouvez instancier un Dataset directement à partir de la table :

ds = Dataset(table)

Vous pouvez accéder à la table PyArrow d'un jeu de données à l'aide de Dataset.data, qui renvoie une MemoryMappedTable ou une InMemoryTable ou une ConcatenationTable, selon l'origine des données Arrow et les opérations appliquées.

Ces objets encapsulent le tableau PyArrow sous-jacent accessible via Dataset.data.table. Ce tableau contient toutes les données du jeu de données, mais il peut également exister une correspondance d’indices dans Dataset._indices qui associe les indices des lignes du jeu de données aux indices des lignes du tableau PyArrow. Cela peut se produire si le jeu de données a été mélangé avec Dataset.shuffle() ou si seul un sous-ensemble des lignes est utilisé (par exemple, après un Dataset.select()).

Dans le cas général, vous pouvez exporter un jeu de données vers une table PyArrow en utilisant table = ds.with_format("arrow")[:].