Transforming a million rows without waiting an hour

The naive version, and the fast one

# slow: one Python call per row
ds = ds.map(lambda r: {"n_chars": len(r["text"])})

# fast: one Python call per thousand rows
ds = ds.map(
    lambda b: {"n_chars": [len(t) for t in b["text"]]},
    batched=True, batch_size=1000,
)

In batched mode the function receives a dict of lists and must return a dict of lists of the same length. Ten to fifty times faster is typical, and the reason is not clever: the per-call overhead of crossing between Arrow and Python dominates, and batching pays it a thousand times less often.

For tokenization the gain is larger still, because the Rust tokenizer parallelises across a batch internally:

def tok_fn(batch):
    return tok(batch["text"], truncation=True, max_length=512)

ds = ds.map(tok_fn, batched=True, remove_columns=["text"])

remove_columns drops the original text once you have the ids. On a large corpus this is the difference between a 30 GB cached file and a 6 GB one.

Parallelism across cores

ds = ds.map(fn, batched=True, num_proc=8)

This forks processes, each mapping a shard, and then concatenates. Two cautions that cost people real time:

• Fast tokenizers already use threads. Combining num_proc with them can oversubscribe the CPU and run slower. The library will warn about TOKENIZERS_PARALLELISM; set it to false when using num_proc, or use one or the other.
• Anything holding a GPU does not fork well. For GPU work, use one process and a large batch, not eight processes fighting over one card.

A batch does not have to stay the same size

This is the feature people miss. A batched map may return more or fewer rows than it received, which makes it the tool for splitting and dropping:

def chunk(batch, size=800, overlap=100):
    out = {"chunk": [], "doc_id": []}
    for doc_id, text in zip(batch["id"], batch["text"]):
        for i in range(0, len(text), size - overlap):
            out["chunk"].append(text[i:i + size])
            out["doc_id"].append(doc_id)
    return out

chunks = docs.map(chunk, batched=True, remove_columns=docs.column_names)

One thousand documents in, forty thousand chunks out, with the source id carried along. remove_columns=docs.column_names is required here, because the old columns no longer have the right length and the library will refuse the mismatch.

filter, and doing it in one pass

ds = ds.filter(lambda b: [len(t) > 50 for t in b["text"]], batched=True, num_proc=4)

Filtering reads everything, so chain your conditions into one predicate rather than calling filter four times. Four filters are four full passes.

with_transform: when you do not want a new file

map materialises a new Arrow file. When the transformation is cheap and varies per epoch (image augmentation, random cropping, dynamic masking) that is wasteful and wrong, because you want different randomness each time.

def augment(batch):
    batch["pixel_values"] = [aug(img) for img in batch["image"]]
    return batch

ds.set_transform(augment)      # applied on access, nothing written

Rule: map for deterministic preprocessing you want cached, set_transform for anything random or cheap.

Making it reproducible

Every map that involves randomness needs a seed passed in, not read from global state, or your cache and your results diverge from each other. And keep the transformation code in a module rather than a notebook cell, because a notebook cell that was edited and re-run is not a record of what produced your data.

Sorting, grouping and joining

sort, class_encode_column and flatten_indices cover most reshaping. There is no join; for that, convert to pandas or Polars, join, and convert back:

import pandas as pd
from datasets import Dataset
merged = Dataset.from_pandas(ds.to_pandas().merge(other_df, on="id"))

That materialises everything in RAM, so it is fine for a few million rows and not for a corpus. For large joins, DuckDB over the parquet files is the next lesson and the right tool.

Do this now

Take one map you have written per row and rewrite it batched. Time both on 50,000 rows. The number you get is the tax you have been paying on every dataset you have ever processed one row at a time.