Use com PyArrow

Este documento é uma introdução rápida ao uso de datasets com PyArrow, com foco especial em como processar datasets usando funções de computação do Arrow e como converter um dataset para o PyArrow ou a partir do PyArrow.

Isso é particularmente útil, pois permite operações rápidas de zero-copy, já que datasets usa PyArrow por baixo dos panos.

Formato do conjunto de dados

Por padrão, os conjuntos de dados retornam objetos Python comuns: inteiros, números de ponto flutuante, strings, listas, etc.

Para obter Tabelas ou Arrays do PyArrow em vez disso, você pode definir o formato do dataset como pyarrow usando 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"
  ]
]

Isso também funciona para objetos IterableDataset obtidos, por exemplo, usando 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]]

Processar dados

As funções do PyArrow são geralmente mais rápidas do que funções python escritas manualmente, e por isso são uma boa opção para otimizar o processamento de dados. Você pode usar as funções de computação do Arrow para processar um dataset em 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]]

Usamos batched=True porque é mais rápido processar lotes de dados no PyArrow do que linha por linha. Também é possível usar batch_size= em map() para definir o tamanho de cada table.

Isso também funciona para IterableDataset.map() e IterableDataset.filter().

Importar ou Exportar do PyArrow

Um Dataset é um wrapper de uma Tabela PyArrow, você pode instanciar um Dataset diretamente da Tabela:

ds = Dataset(table)

Você pode acessar a Tabela PyArrow de um dataset usando Dataset.data, que retorna uma MemoryMappedTable ou uma InMemoryTable ou uma ConcatenationTable, dependendo da origem dos dados Arrow e das operações aplicadas.

Esses objetos encapsulam a tabela PyArrow subjacente, acessível em Dataset.data.table. Essa tabela contém todos os dados do conjunto de dados, mas também pode haver um mapeamento de índices em Dataset._indices, que mapeia os índices das linhas do conjunto de dados para os índices das linhas da Tabela PyArrow. Isso pode ocorrer se o conjunto de dados tiver sido embaralhado com Dataset.shuffle() ou se apenas um subconjunto das linhas for usado (por exemplo, após um Dataset.select()).

No caso geral, você pode exportar um dataset para uma Tabela PyArrow usando table = ds.with_format("arrow")[:].