Datasets documentation
与 Polars 配合使用
与 Polars 配合使用
本文档简要介绍了如何使用 datasets 与 Polars 配合，重点说明如何使用 Polars 函数处理数据集，以及如何在数据集与 Polars 之间进行转换。
这一点尤其有用，因为它支持快速的零拷贝操作，因为 datasets 和 Polars 底层都使用 Arrow。
数据集格式
默认情况下，数据集返回的是普通的 Python 对象：整数、浮点数、字符串、列表等。
To get Polars DataFrames or Series instead, you can set the format of the dataset to polars using 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"
]对于例如使用 load_dataset(..., streaming=True) 获取的 IterableDataset 对象，此方法同样适用：
>>> 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 │
└───────┴───────┘处理数据
Polars 函数通常比常规手写 Python 函数更快，因此是优化数据处理的理想选择。您可以在 Dataset.map() 或 Dataset.filter() 中使用 Polars 函数来处理数据集：
>>> 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 │
└───────┴───────┴───────┘我们使用 batched=True，因为在 Polars 中按批次处理数据比逐行处理更快。也可以在 map() 中使用 batch_size= 来设置每个 df 的大小。
这对 IterableDataset.map() 和 IterableDataset.filter() 同样适用。
示例：数据提取
Polars 提供了适用于各种数据类型（如字符串、浮点数、整数等）的众多函数。您可以在此处找到完整列表。这些函数使用 Rust 编写，并对数据批次进行处理，从而实现快速的数据处理。
下面是一个示例，展示了使用 Polars 代替普通 Python 函数从 LLM 推理数据集中提取解决方案时，速度提升了 5 倍：
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: 2s从 Polars 导入或导出
要从 Polars 导入数据，您可以使用 Dataset.from_polars()：
ds = Dataset.from_polars(df)
您可以使用 Dataset.to_polars() 将 Dataset 导出为 Polars DataFrame：
df = Dataset.to_polars(ds)