資料集文件
與 Polars 搭配使用
並獲得增強的文件體驗
開始使用
與 Polars 搭配使用
本文件簡要介紹如何將 datasets 與 Polars 結合使用，特別著重於如何利用 Polars 函式處理資料集，以及如何在資料集與 Polars 之間進行格式轉換。
這非常實用，因為由於 datasets 和 Polars 底層皆使用 Arrow，因此能實現快速的零複製（zero-copy）操作。
資料集格式
預設情況下，資料集會傳回一般的 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() 將資料集匯出為 Polars DataFrame。
df = Dataset.to_polars(ds)