与 Polars 一起使用

本文档是对将 datasets 与 Polars 结合使用的快速介绍，特别关注如何使用 Polars 函数处理数据集，以及如何将数据集转换为 Polars 或从 Polars 转换数据集。

这特别有用，因为它允许快速的零拷贝操作，因为 datasets 和 Polars 在底层都使用 Arrow。

数据集格式

默认情况下，datasets 返回常规 Python 对象：整数、浮点数、字符串、列表等。

要获取 Polars DataFrame 或 Series，你可以使用 Dataset.with_format() 将数据集的格式设置为 polars。

>>> 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.from_polars(ds)

< > 在 GitHub 上更新