与 Polars 结合使用

本文档简要介绍了如何将 datasets 与 Polars 结合使用，特别关注如何使用 Polars 函数处理数据集，以及如何将数据集与 Polars 相互转换。

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

数据集格式

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

要获取 Polars DataFrames 或 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 从一个大语言模型（LLM）推理数据集中提取解决方案时，速度比常规 Python 函数快了 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)

< > 在 GitHub 上更新