map, batched and the cache
map applies your function to every row, batched=True hands it whole chunks at once, and a fingerprint of your function decides whether the work is redone or read back from disk.
- 9 min read
- 3 reading levels
- Published
Read these first
On this page 5
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
map runs your preparation function over every row of a dataset, and it saves the result so the same work never happens twice.
Think of prepping vegetables for a week of cooking. You chop everything on Sunday and put it in boxes in the fridge. On Wednesday you open a box instead of picking up the knife again. The chopping happened once; the meals reuse it.
map is that Sunday chopping for training data. Turning text into numbers is slow, and you would otherwise pay for it before every single training run.
Why it exists
Models eat numbers, not sentences. Something must convert every row, and that conversion can take minutes on a big dataset.
Doing it inside the training loop means paying the cost every epoch. Doing it once and saving the result means paying once. map is where that happens, and the saved boxes are its cache.
There is a second trick. Handing the chopper one carrot at a time wastes most of its speed. Handing over a whole basket lets it work in bulk. That is what batched=True does.
How it works
raw rows your function new columns
"cold biryani" ──map(encode)──→ [101, 3147, 12170, ...]
batched=False: one row per call ← lots of small calls, slow
batched=True : many rows per call ← few big calls, fast
second run: same data + same function → read the saved boxes, do nothingThe "same function" check is the part that surprises people. The library takes a fingerprint of your function and the data. Change either one and the work is redone. Change neither and it is skipped.
A real example you have seen
Your phone's photo search. The first time it offers to find "dogs", it spends a night working through your gallery. After that, searching is instant — because the hard part was done once and stored.
Remember this
mapapplies your function to a dataset and returns a new one.batched=Truehands your function many rows per call, which is much faster.- Results are cached on disk; an unchanged function and dataset means no work at all.
What to learn next
- Data collators — the per-batch work that deliberately stays out of
map. - The Trainer API — where the mapped dataset finally gets used.
- Is the GPU waiting for data? — how to tell whether preprocessing is your bottleneck.
Developer — Code and libraries.
Setup
pip install datasets transformersTested with datasets 5.0 and transformers 5.6. Tokenizer files for google/bert_uncased_L-2_H-128_A-2 are about 18 MB. Everything runs on CPU.
What "batched" actually hands your function
from datasets import Dataset
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
ds = Dataset.from_dict({"text": ["cold biryani", "great filter coffee",
"slow service", "best vada pav"],
"stars": [1, 5, 2, 5]})
calls = {"n": 0}
def encode(rows):
calls["n"] += 1
print(f" call {calls['n']}: {len(rows['text'])} texts -> {rows['text']}")
return tok(rows["text"], truncation=True, max_length=16)
enc = ds.map(encode, batched=True, batch_size=2, remove_columns=["text"])
print("total calls:", calls["n"], "for", ds.num_rows, "rows")
print(enc)
print("row 0:", enc[0]) call 1: 2 texts -> ['cold biryani', 'great filter coffee']
call 2: 2 texts -> ['slow service', 'best vada pav']
total calls: 2 for 4 rows
Dataset({
features: ['stars', 'input_ids', 'token_type_ids', 'attention_mask'],
num_rows: 4
})
row 0: {'stars': 1, 'input_ids': [101, 3147, 12170, 20444, 3490, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1]}With batched=True your function receives a dict of lists, not a row. rows["text"] is a list of two strings, and the tokenizer handles lists natively — which is the whole reason batching is faster.
What it costs, and what caching saves
import os, csv, time
os.environ["HF_DATASETS_CACHE"] = os.path.abspath("ds_cache") # keep this demo self-contained
from datasets import load_dataset
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
with open("big_reviews.csv", "w", newline="", encoding="utf-8") as f:
w = csv.writer(f); w.writerow(["text", "stars"])
for i in range(20000):
w.writerow([f"review {i}: the dosa was fine and the coffee was hot", i % 5 + 1])
ds = load_dataset("csv", data_files="big_reviews.csv", split="train")
def encode_one(row): return tok(row["text"], truncation=True, max_length=32)
def encode_batch(rows): return tok(rows["text"], truncation=True, max_length=32)
t = time.time(); ds.map(encode_one); one = time.time() - t
t = time.time(); a = ds.map(encode_batch, batched=True, batch_size=1000); many = time.time() - t
t = time.time(); b = ds.map(encode_batch, batched=True, batch_size=1000); again = time.time() - t
print(f"row at a time : {one:.2f}s")
print(f"batched=True : {many:.2f}s")
print(f"run it again : {again:.2f}s")
print("same fingerprint:", a._fingerprint == b._fingerprint, a._fingerprint)
print("cache file:", os.path.basename(a.cache_files[0]["filename"]))row at a time : 1.44s batched=True : 0.25s run it again : 0.00s same fingerprint: True c06578902df43370 cache file: cache-c06578902df43370.arrow
Timings depend on your CPU; the ratios are the point. Roughly six times faster batched, and free on the second run. The fingerprint hash will differ on your machine.
The walkthrough
The cache file is named after the fingerprint. cache-c06578902df43370.arrow is not a coincidence — the fingerprint is the filename. The library hashes the input dataset's fingerprint together with your function's bytecode and arguments. Same hash, file already there, nothing to do.
remove_columns is about memory, not tidiness. Without it the raw text column is copied into the new Arrow file alongside the token ids. On a large corpus that doubles the file. Also, the Trainer rejects string columns it cannot turn into tensors — so the string usually has to go anyway.
Do not pad inside map. It is tempting to pass padding="max_length" here. That freezes every row to the same length forever and wastes compute on every short one. Pad per batch instead, at the last moment — the job of data collators.
num_proc parallelises across cores. ds.map(encode_batch, batched=True, num_proc=4) forks worker processes. It is a large win for slow Python functions and a loss for fast-tokenizer work, which already uses every core inside Rust. Measure before adopting it. On Windows and macOS, code calling num_proc must sit inside if __name__ == "__main__": or the processes re-import your script and fail.
set_transform is the lazy alternative. It applies your function on access instead of writing a new table. Correct for heavy image augmentation that should differ every epoch; wrong for tokenization, which you want computed once.
Common mistakes
Getting a stale cache after editing your function. Fingerprinting reads your function's code, so real edits do produce a new hash. Values captured from outside — a max_length read from a global, a tokenizer swapped in place — can slip through. When output looks impossibly wrong, pass load_from_cache_file=False once to confirm, then find what your function was reading from outside itself.
A function that returns the wrong shape. With batched=True, every returned list must be as long as the batch you were given, or your columns misalign with the rest of the table. You can deliberately return a different length — for chunking long documents into windows — but only when every returned column shares that new length.
Silently keeping the old column. map adds columns; it does not replace them. Returning {"text": cleaned} overwrites text, while returning {"clean_text": cleaned} leaves both. Print column_names after every map until this stops surprising you.
A cache directory that grows forever. Every distinct function writes another Arrow file. ds.cleanup_cache_files() deletes the ones for that dataset. Watch it on a laptop, especially with big image datasets.
Unhashable arguments. A function carrying an object the fingerprinter cannot hash makes datasets fall back to a random fingerprint. Then nothing ever caches, and every run redoes the work. The warning about a non-hashable object is worth reading rather than ignoring.
Try it yourself
Change max_length from 32 to 16 and rerun timing.py. Does the second run still take 0.00s, and what does the fingerprint become? Then set num_proc=4 on the 20,000-row map and time it — on fast-tokenizer work, is it faster or slower than the single-process version on your machine?
What to learn next
- Data collators — the per-batch work that deliberately stays out of
map. - The Trainer API — where the mapped dataset finally gets used.
- Is the GPU waiting for data? — how to tell whether preprocessing is your bottleneck.
Researcher — Mathematics and papers.
Fingerprinting as pure-function bookkeeping
Each Dataset carries a _fingerprint. Transformations are decorated so the new fingerprint is hash(parent_fingerprint, transform_name, hashed_kwargs), where callables are hashed by a custom Pickler that walks __code__, closure cells, default arguments and referenced globals. The scheme assumes transforms are pure functions of their inputs — a lineage hash, in the same spirit as build systems and dbt models.
Where the assumption leaks is worth naming precisely, because it is where stale caches come from:
- A closure over a mutable object hashes the object's current state at map time. Mutating it in place afterwards is invisible.
- Non-determinism inside the function (randomness, wall-clock, network) is not represented at all; the first result is cached as if canonical.
- Objects the pickler cannot serialise trigger a random fingerprint, disabling caching silently except for a warning.
datasets.disable_caching() turns the whole mechanism off for a session, which is the honest setting for randomised augmentation pipelines.
Batching, and the shape contract
batched=True slices the table into contiguous chunks of batch_size and passes a dict of columns. Two properties matter. Chunks are contiguous in Arrow, so the slice is a zero-copy view. And the returned column lengths need only agree with each other, not with the input — so one input row may yield many output rows.
That is the mechanism behind sliding-window question answering, where return_overflowing_tokens=True produces several windows per document and overflow_to_sample_mapping records which document each came from; the offsets in fast tokenizers then map predictions back. It is also how packing is implemented for causal-LM pretraining: concatenate a batch's token streams, then re-chunk to a fixed block length.
The speedup itself is not vectorisation in the numeric sense. It is amortisation of per-call Python and Arrow overhead, plus — for fast tokenizers — Rayon-parallel encoding across cores inside the Rust tokenizers library. This is exactly why stacking num_proc on top often loses: the cores are already busy, and you have added process forks and inter-process serialisation of Arrow batches.
Where the cost actually belongs
The general principle is to hoist deterministic, expensive, input-independent work out of the training loop, and to leave stochastic work inside it. Tokenization is deterministic given the tokenizer, so it is cached. Padding depends on batch composition, so it belongs in the collator. Augmentation is intentionally random per epoch, so it belongs in set_transform or in a PyTorch Dataset.__getitem__ — see torchvision transforms v2. Getting an item on the wrong side of that line is the most common preprocessing performance bug in practice.
What to learn next
- Data collators — the per-batch work that deliberately stays out of
map. - The Trainer API — where the mapped dataset finally gets used.
- Is the GPU waiting for data? — how to tell whether preprocessing is your bottleneck.