Datasets documentation
Uso con Polars
Uso con Polars
Este documento es una introducción rápida al uso de datasets con Polars, con un enfoque particular en cómo procesar conjuntos de datos mediante funciones de Polars y cómo convertir un conjunto de datos a Polars o desde Polars.
Esto resulta especialmente útil, ya que permite operaciones rápidas de copia cero, ya que tanto datasets como Polars utilizan Arrow internamente.
Formato del conjunto de datos
Por defecto, los conjuntos de datos devuelven objetos de Python estándar: enteros, flotantes, cadenas, listas, etc.
Para obtener DataFrames o Series de Polars en su lugar, puede establecer el formato del conjunto de datos en 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"
]Esto también funciona para objetos IterableDataset obtenidos, por ejemplo, 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 │
└───────┴───────┘Procesar datos
Las funciones de Polars son generalmente más rápidas que las funciones de python escritas a mano, por lo que son una buena opción para optimizar el procesamiento de datos. Puedes usar funciones de Polars para procesar un conjunto de datos en Dataset.map() o 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 es más rápido procesar lotes de datos en Polars que fila por fila. También es posible usar batch_size= en map() para establecer el tamaño de cada df.
Esto también funciona para IterableDataset.map() y IterableDataset.filter().
Ejemplo: extracción de datos
Polars ofrece muchas funciones disponibles para cualquier tipo de datos: cadenas, números de coma flotante, enteros, etc. Puedes encontrar la lista completa aquí. Estas funciones están escritas en Rust y se ejecutan en lotes de datos, lo que permite un procesamiento de datos rápido.
A continuación se muestra un ejemplo que ilustra una aceleración de 5x al usar Polars en lugar de una función de python regular para extraer soluciones de un conjunto de datos de razonamiento de 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: 2sImportar o exportar desde Polars
Para importar datos desde Polars, puedes usar Dataset.from_polars():
ds = Dataset.from_polars(df)
Y puedes usar Dataset.to_polars() para exportar un Dataset a un DataFrame de Polars:
df = Dataset.to_polars(ds)