Datasets documentation
Utiliser avec Polars
Utiliser avec Polars
Ce document est une introduction rapide à l'utilisation de datasets avec Polars, avec un accent particulier sur la manière de traiter les datasets à l'aide de fonctions Polars, et sur la conversion d'un dataset vers Polars ou depuis Polars.
C'est particulièrement utile car cela permet des opérations rapides sans copie, car datasets et Polars utilisent tous deux Arrow 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 DataFrames ou des Series Polars à la place, vous pouvez définir le format du jeu de données sur polars à 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("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"
]Cela fonctionne également pour les objets IterableDataset obtenus, par exemple, à l'aide de 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 │
└───────┴───────┘Traiter les données
Les fonctions Polars 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 Polars pour traiter un jeu de données dans 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 │
└───────┴───────┴───────┘Nous utilisons batched=True car il est plus rapide de traiter des lots de données dans Polars que ligne par ligne. Il est également possible d’utiliser batch_size= dans map() pour définir la taille de chaque df.
Cela fonctionne également pour IterableDataset.map() et IterableDataset.filter().
Exemple : extraction de données
De nombreuses fonctions sont disponibles dans Polars et pour tout type de données : chaînes de caractères, flottants, entiers, etc. Vous pouvez consulter la liste complète ici. Ces fonctions sont écrites en Rust et s'exécutent sur des lots de données, ce qui permet un traitement rapide des données.
Voici un exemple montrant un gain de vitesse de 5x en utilisant Polars plutôt qu'une fonction python classique pour extraire les solutions d'un jeu de données de raisonnement d'un 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: 2sImporter ou exporter depuis Polars
Pour importer des données depuis Polars, vous pouvez utiliser Dataset.from_polars() :
ds = Dataset.from_polars(df)
Vous pouvez utiliser Dataset.to_polars() pour exporter un Dataset vers un DataFrame Polars :
df = Dataset.to_polars(ds)