tf.data input pipelines

The data ingestion framework. Spans tensorflow/python/data/ (Python API), tensorflow/core/data/ (runtime), and tensorflow/core/kernels/data/ (op kernels).

Purpose

• Build high-throughput input pipelines that overlap data loading and preprocessing with model training.
• Be backed by stateful kernels (Dataset, Iterator) implemented in C++ for performance.
• Be composable through Python operators (map, batch, shuffle, prefetch, interleave, cache, …).
• Support distributed input (per-replica datasets, sharding) via tf.distribute.

Directory layout

tensorflow/python/data/
├── ops/                        # Dataset class + every transformation
│   ├── dataset_ops.py          # Core: tf.data.Dataset
│   ├── iterator_ops.py
│   ├── map_op.py
│   ├── batch_op.py
│   ├── shuffle_op.py
│   ├── prefetch_op.py
│   ├── interleave_op.py
│   ├── cache_op.py
│   ├── tf_record_op.py
│   └── ... (one file per transformation)
├── experimental/                # Less-stable transformations
├── benchmarks/
└── kernel_tests/

tensorflow/core/data/
├── dataset_utils.{h,cc}
├── snapshot_utils.{h,cc}        # tf.data snapshot
├── service/                     # tf.data service: distributed input
├── ...

tensorflow/core/kernels/data/
├── batch_dataset_op.{h,cc}
├── cache_dataset_ops.{h,cc}
├── concatenate_dataset_op.{h,cc}
├── filter_dataset_op.{h,cc}
├── interleave_dataset_op.{h,cc}
├── iterator_ops.{h,cc}          # MakeIterator, IteratorGetNext, ...
├── map_dataset_op.{h,cc}
├── parallel_interleave_dataset_op.{h,cc}
├── parallel_map_dataset_op.{h,cc}
├── prefetch_dataset_op.{h,cc}
├── repeat_dataset_op.{h,cc}
├── shuffle_dataset_op.{h,cc}
├── window_dataset_op.{h,cc}
├── tf_record_dataset_op.{h,cc}
└── ...

Key abstractions

Type / class	File	Purpose
tf.data.Dataset	tensorflow/python/data/ops/dataset_ops.py	The user-facing dataset.
Iterator (Python)	tensorflow/python/data/ops/iterator_ops.py	Yields elements in eager / function code.
DatasetBase (C++)	tensorflow/core/framework/dataset.h	Base class for runtime datasets.
IteratorBase (C++)	tensorflow/core/framework/dataset.h	Stateful iterator that produces elements.
IteratorResource	tensorflow/core/kernels/data/iterator_ops.h	Resource that owns an IteratorBase.
tf.data.Options	tensorflow/python/data/ops/options.py	Pipeline-level options (autotune, sharding).
OptimizeDataset op	tensorflow/core/kernels/data/optimize_dataset_op.cc	Inserts pipeline optimizations (map_and_batch_fusion, etc.).

How a pipeline executes

graph LR
    Py["ds = tf.data.TFRecordDataset(...).map(parse).batch(32).prefetch(2)"]
    Build[Python builds a tree of Dataset variants]
    Make[MakeIterator op creates IteratorResource]
    Get[IteratorGetNext op reads one element]
    Threads[Background prefetch/parallel_map threads]
    Trainer[Trainer loop / model.fit]

    Py --> Build
    Build --> Make
    Make --> Get
    Threads --> Get
    Get --> Trainer

1. The user composes a tf.data.Dataset in Python. Each transformation builds a variant tensor representing the dataset (basically a serialised "how to produce elements" recipe).
2. Dataset.__iter__ (eager) or iter(dataset) constructs an IteratorResource via the MakeIterator op kernel.
3. Each next(...) call invokes IteratorGetNext, which the underlying iterator implementation (a C++ class registered for that dataset op) services. Iterators with internal threads (PrefetchDataset, ParallelMapDataset, ParallelInterleaveDataset) maintain background workers that fill an internal buffer.
4. The trainer pulls batches from the iterator inside a @tf.function-wrapped step.

Optimization

tf.data has its own pipeline optimizer that runs over the dataset graph before it executes:

• map_and_batch_fusion — fuses Dataset.map().batch() into a single op when the map function is stateless.
• noop_elimination, shuffle_and_repeat_fusion, parallel_batch, inject_prefetch, autotune_buffer_sizes.
• The optimizer entry is OptimizeDataset (tensorflow/core/kernels/data/optimize_dataset_op.cc).
• Auto-tuning (tf.data.AUTOTUNE) uses a runtime model in tensorflow/core/framework/model.{h,cc} to pick parallelism and buffer sizes dynamically.

Snapshot and caching

• Dataset.cache() materialises a pass through to memory or disk; backed by cache_dataset_ops.cc.
• Dataset.snapshot() writes an on-disk snapshot for reuse across runs (tensorflow/core/data/snapshot_utils.cc).

tf.data service

tensorflow/core/data/service/ is a separate distributed input service. The user runs a tf.data dispatcher + workers on dedicated machines, and a Python client distributes input across them. Useful when the model trainer is GPU-bound and CPU-side input would otherwise become the bottleneck.

Distribution

tf.distribute integrates by wrapping a Dataset into a DistributedDataset (tensorflow/python/distribute/input_lib.py). Sharding is picked per-strategy: MirroredStrategy shards across replicas in a single host; MultiWorkerMirroredStrategy shards across hosts.

Integration points

• Variant tensors — every dataset is a DT_VARIANT tensor referencing a DatasetBase C++ object.
• Resource model — iterators are Resources managed by a ResourceMgr.
• Checkpointing — iterators can save/restore via tf.train.Checkpoint; the IteratorBase::Save/Restore interface and tf.data checkpoint ops are recently security-hardened (see, e.g., 2026 commits like [tf.data Security] Validate checkpoint values.).
• @tf.function — iterators are captured into traced graphs; Dataset.reduce is the canonical way to write a graph-mode loop over a dataset.

Entry points for modification

• New transformation: add a Python class in tensorflow/python/data/ops/<name>_op.py, an OpDef in tensorflow/core/ops/dataset_ops.cc, and a kernel under tensorflow/core/kernels/data/. The kernel implements DatasetBase::Iterator.
• New optimization: add an MLIR-style pass under tensorflow/core/grappler/optimizers/data/ and wire it into OptimizeDataset.
• tf.data service changes: tensorflow/core/data/service/.

Related

• systems/core-runtime — variant tensors, resource manager.
• features/distribution-strategy — distributed input.