Datasets documentation
Use com Polars
Use com Polars
Este documento é uma introdução rápida ao uso de datasets com Polars, com foco especial em como processar datasets usando funções do Polars e como converter um dataset para o Polars ou a partir do Polars.
Isso é particularmente útil, pois permite operações rápidas de zero-copy, já que tanto datasets quanto o Polars usam Arrow 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 DataFrames ou Series do Polars em vez disso, você pode definir o formato do dataset como polars 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("polars")
>>> ds[0] # pl.DataFrame
shape: (1, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
└───────┴───────┘
>>> ds[:2] # pl.DataFrame
shape: (2, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
│ b ┆ 0.0 │
└───────┴───────┘
>>> ds["data"] # pl.Series
shape: (4,)
Series: 'col_0' [str]
[
"a"
"b"
"c"
"d"
]Isso também funciona para objetos IterableDataset obtidos, por exemplo, usando load_dataset(..., streaming=True):
>>> ds = ds.with_format("polars")
>>> for df in ds.iter(batch_size=2):
... print(df)
... break
shape: (2, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
│ b ┆ 0.0 │
└───────┴───────┘Processar dados
As funções do Polars são geralmente mais rápidas do que funções python escritas manualmente, e, portanto, são uma boa opção para otimizar o processamento de dados. Você pode usar funções do Polars para processar um conjunto de dados em Dataset.map() ou Dataset.filter():
>>> import polars as pl
>>> 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("polars")
>>> ds = ds.map(lambda df: df.with_columns(pl.col("col_1").add(1).alias("col_2")), batched=True)
>>> ds[:2]
shape: (2, 3)
┌───────┬───────┬───────┐
│ col_0 ┆ col_1 ┆ col_2 │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═══════╪═══════╪═══════╡
│ a ┆ 0.0 ┆ 1.0 │
│ b ┆ 0.0 ┆ 1.0 │
└───────┴───────┴───────┘
>>> ds = ds.filter(lambda df: df["col_0"] == "b", batched=True)
>>> ds[0]
shape: (1, 3)
┌───────┬───────┬───────┐
│ col_0 ┆ col_1 ┆ col_2 │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═══════╪═══════╪═══════╡
│ b ┆ 0.0 ┆ 1.0 │
└───────┴───────┴───────┘Usamos batched=True porque é mais rápido processar lotes de dados no Polars do que linha por linha. Também é possível usar batch_size= em map() para definir o tamanho de cada df.
Isso também funciona para IterableDataset.map() e IterableDataset.filter().
Exemplo: extração de dados
Muitas funções estão disponíveis no Polars e para qualquer tipo de dados: string, floats, integers, etc. Você pode encontrar a lista completa aqui. Essas funções são escritas em Rust e executam em lotes de dados, o que permite um processamento de dados rápido.
Aqui está um exemplo que mostra um aumento de 5x na velocidade ao usar Polars em vez de uma função python comum para extrair soluções de um conjunto de dados de raciocínio de LLM:
from datasets import load_dataset
ds = load_dataset("ServiceNow-AI/R1-Distill-SFT", "v0", split="train")
# Using a regular python function
pattern = re.compile("boxed\\{(.*)\\}")
result_ds = ds.map(lambda x: {"value_solution": m.group(1) if (m:=pattern.search(x["solution"])) else None})
# Time: 10s
# Using a Polars function
expr = pl.col("solution").str.extract("boxed\\{(.*)\\}").alias("value_solution")
result_ds = ds.with_format("polars").map(lambda df: df.with_columns(expr), batched=True)
# Time: 2sImportar ou Exportar do Polars
Para importar dados do Polars, você pode usar Dataset.from_polars():
ds = Dataset.from_polars(df)
E você pode usar Dataset.to_polars() para exportar um Dataset para um DataFrame do Polars:
df = Dataset.to_polars(ds)