Datasets documentation
Polars との併用
Polars との併用
このドキュメントは、datasetsをPolarsと組み合わせて使用する際の簡単な入門ガイドです。特に、Polars関数を用いたデータセットの処理方法、およびデータセットをPolarsへの変換とPolarsからの変換方法に焦点を当てています。
datasetsとPolarsはどちらも内部でArrowを使用しているため、ゼロコピー操作を高速に行える点で特に有用です。
データセットの形式
デフォルトでは、datasetsは通常の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"
]IterableDatasetオブジェクト（例えば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 │
└───────┴───────┘データを処理する
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 │
└───────┴───────┴───────┘Polarsでは行単位ではなくバッチ単位でデータを処理する方が高速なため、batched=Trueを使用します。また、map()内でbatch_size=を指定することで、各dfのサイズを設定することも可能です。
これはIterableDataset.map()およびIterableDataset.filter()にも適用できます。
例: データ抽出
Polarsでは、文字列、浮動小数点数、整数など、あらゆるデータ型に対して多数の関数が利用可能です。関数の完全なリストはこちらで確認できます。これらの関数はRustで記述されており、データバッチ単位で実行されるため、高速なデータ処理を実現します。
通常のPython関数ではなくPolarsを使用して、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: 2sPolarsからのインポートまたはエクスポート
Polarsからデータをインポートするには、Dataset.from_polars()を使用できます:
ds = Dataset.from_polars(df)
Dataset.to_polars() を使用して、Dataset を Polars DataFrame にエクスポートできます:
df = Dataset.to_polars(ds)