Dataloader

rerun.experimental.dataloader

PyTorch Datasets for training on data from the Rerun catalog.

IndexValue module-attribute

IndexValue: TypeAlias = int | datetime64 | timedelta64

A concrete index value on a timeline.

A plain int for integer timelines, a datetime64[ns] for timestamp timelines, and a timedelta64[ns] for duration timelines.

BlockShuffle dataclass

Bases: ShuffleStrategy

Shuffle fetch-sized blocks of consecutive samples instead of individual samples.

Each block is one fetch's worth of consecutive, segment-local samples. The block order is shuffled, but samples keep their natural order within a block, so every fetch still reads one contiguous span: stored data is read about once per epoch instead of once per fetch, and decoders reuse their cache across a block. Set buffer_size to decorrelate batches at emission time.

PARAMETER	DESCRIPTION
buffer_size	Size of the post-decode buffer that randomizes emission order without changing the fetch order. None (the default) emits in fetch order. This is the only strategy that takes one — see emission_buffer. Warning The buffer holds up to buffer_size decoded samples per DataLoader worker, as your decoders produce them — full-resolution frames for video fields. Budget buffer_size * bytes_per_sample * num_workers; a few thousand video samples per worker is tens of gigabytes. TYPE: int | None DEFAULT: None
min_fill	Samples buffered before emission starts. Defaults to buffer_size // 2, which is also how long the first sample is delayed; lower it to shorten that warm-up at the cost of less mixing over the first few batches. TYPE: int | None DEFAULT: None

ColumnDecoder

Bases: ABC, Generic[_DecodedT_co]

Base class for column decoders.

Subclasses convert raw Arrow data into decoded values. The pipeline calls decode once per field with every requested sample of a fetch block, so decoders can amortize work across samples (one vectorized gather for numeric data, one codec pass per GOP for video). Decoders that only care about one sample at a time can simply loop over the requests and decode each FieldBatch.take_decode_rows window.

fill_latest_at property

fill_latest_at: bool

Whether this column's prefetch read latest-at-fills empty grid slots.

True for stateless columns (images, scalars): each grid slot wants the most recent value snapped from the real rows. Keyframe-aware fields use exact range queries instead, regardless of this value. The read is partitioned by this flag so it stays a global query argument per group rather than a per-column one.

decode abstractmethod

def decode(
    batch: FieldBatch, requests: Sequence[DecodeRequest]
) -> Sequence[_DecodedT_co | None]

Decode all requests against batch, returning one entry per request.

requests arrive in row order: grouped by segment, and ascending by index value within a segment. Implementations may process them in any internal order (e.g. grouped by GOP), but the result must align 1:1 with the input: result[i] is requests[i]'s decoded value, or None to signal data missing for that sample.

prior_keyframe_path

def prior_keyframe_path(field_path: str) -> str | None

Sibling column whose non-null rows mark a re-entrant keyframe, or None.

Override on decoders that need the prefetch window anchored at the prior keyframe (compressed video). Default returns None.

DataSource dataclass

An immutable reference to a dataset with an optional segment filter.

PARAMETER	DESCRIPTION
dataset	The remote dataset to read from. TYPE: DatasetEntry
segments	Optional list of segment IDs to restrict to. TYPE: list[str] | None DEFAULT: None

filter_segments

def filter_segments(segment_ids: list[str]) -> DataSource

Return a new DataSource narrowed to segment_ids.

DecodeRequest dataclass

One sample's decode request within a field's batch.

The pipeline has already resolved the index-value window this sample needs into batch rows, so a decoder consumes row indices rather than searching for them.

decode_row_indices instance-attribute

decode_row_indices: tuple[int, ...]

The batch rows holding all data needed to decode this sample.

For compressed video this includes the intermediate frames needed to decode the requested output, even when the field has no explicit window.

index_value instance-attribute

index_value: IndexValue

The typed target index value of this sample.

output_row_indices instance-attribute

output_row_indices: tuple[int, ...]

Physical batch rows whose decoded values form this request's output.

Always non-empty: preparation omits unresolved requests before invoking a decoder.

sample_position instance-attribute

sample_position: int

Position where this request's result belongs in the decoded fetch block.

segment_id instance-attribute

segment_id: str

Segment this sample comes from. Index values are only comparable within one segment.

starts_at_keyframe instance-attribute

starts_at_keyframe: bool

Whether the window's first row is known to be a keyframe a decoder may start from.

True when the pipeline found a prior keyframe at the first decode row; false when no prior keyframe exists for the requested output.

Field dataclass

Declarative spec for one field of a training sample.

Note

This API is provisional and will be improved, expect the surface to change.

PARAMETER	DESCRIPTION
path	entity_path:Archetype:component triple identifying the source column (e.g. "/camera:EncodedImage:blob"). TYPE: str
decode	A ColumnDecoder that turns the Arrow column into a training value. TYPE: ColumnDecoder[DecodedValue]
select	Optional jq-like Selector applied client-side to the Arrow column before decode. Used for nested struct/list access. The server-side projection is unaffected. Field( path="/agent:ListOfStructs:animals", select=Selector(".[0].dog"), decode=NumericDecoder(), ) TYPE: Selector | None DEFAULT: None
window	Optional explicit offsets of the values to return relative to the current index value. Integer timelines require integral index-step offsets. Timestamp and duration timelines use seconds, which are converted to nanoseconds internally. Unlike FixedRateSampling, these offsets do not define or interpolate a grid. For example, (-2.5, 0.0) on a timestamp timeline requests exactly the values at 2.5 seconds before the current sample and at the current sample. An RGB compressed-video window yields a [T, C, H, W] frame stack; VideoFrameDecoder(output_format="yuv420p") instead yields a Yuv420Frame. Both are bootstrapped from the keyframe preceding the earliest output. TYPE: tuple[int | float, ...] | None DEFAULT: None
max_staleness	Optional maximum age of the data backing a sample, using the same unit convention as window: integral index steps for integer timelines and seconds for timestamp or duration timelines. When set, a required sample is dropped if the nearest value at or before a queried point is older than this. None (the default) applies no staleness limit. Enforced only during manifest construction, not by the streaming dataloader. For compressed video, manifests conservatively measure age from the latest prior keyframe because sparse keyframe metadata does not expose non-keyframe timestamps. A sample may therefore be dropped even when a fresher non-keyframe exists. TYPE: int | float | None DEFAULT: None

fill_latest_at property

fill_latest_at: bool

Whether server queries for this field use latest-at filling.

prior_keyframe_path property

prior_keyframe_path: str | None

Component path containing the keyframe markers required by this field's decoder.

to_recipe

def to_recipe() -> dict[str, Any]

A JSON-serializable snapshot of this field's spec, for a manifest's provenance header.

Kept on Field so it stays in sync as the spec evolves. decode / select are captured via repr — a human-readable record, not a round-trippable form.

FieldBatch dataclass

One field's rows for a whole fetch block, across every segment it touches.

Holds the entire fetched window of the field's column, so a decoder can process every sample of the block in one vectorized pass instead of one call per sample. A shuffled sampler puts almost every sample in a different segment, so batching per segment would collapse to a single sample per call.

Rows are ordered by segment, and ascending by index value within a segment, so a decoder walking requests in order walks the column forwards. Index values are only comparable inside a segment, which is why locating a sample's rows is the pipeline's job: see DecodeRequest.

column instance-attribute

column: Array

A Rerun component column, with one outer Arrow list per timeline row.

is_windowed class-attribute instance-attribute

is_windowed: bool = False

Whether decoded outputs keep a leading time axis, including a one-value window.

take_decode_rows

def take_decode_rows(request: DecodeRequest) -> Array

The request's decode rows, with select applied.

take_output_rows

def take_output_rows(request: DecodeRequest) -> Array

The rows requested for output, preserving repeats and applying select.

FixedRateSampling dataclass

Sample timestamp or duration timelines at a fixed nominal rate.

Indices are drawn on an algebraic grid seg.index_start + k * ns_per_sample. The server's fill_latest_at absorbs any drift from real-row positions.

ImageDecoder

Bases: ColumnDecoder[Tensor]

Decode encoded-image blobs (JPEG/PNG) to [C, H, W] uint8 tensors.

fill_latest_at property

fill_latest_at: bool

Whether this column's prefetch read latest-at-fills empty grid slots.

True for stateless columns (images, scalars): each grid slot wants the most recent value snapped from the real rows. Keyframe-aware fields use exact range queries instead, regardless of this value. The read is partitioned by this flag so it stays a global query argument per group rather than a per-column one.

prior_keyframe_path

def prior_keyframe_path(field_path: str) -> str | None

Sibling column whose non-null rows mark a re-entrant keyframe, or None.

Override on decoders that need the prefetch window anchored at the prior keyframe (compressed video). Default returns None.

Manifest

A description of the exact sampling order of data for one epoch of a training run. This serves the purpose of making training runs reproducible and resumable.

A manifest is created by generating one from a source with Manifest.generate, or by loading an existing one with Manifest.from_parquet. Persist it with Manifest.write_parquet; a manifest is always backed by a parquet file on the read path so a DataLoader worker never holds the whole table in RAM.

Note

This API is provisional and will be improved, expect the surface to change.

metadata property

metadata: ManifestMeta

Decoded metadata header.

num_rows property

num_rows: int

Total number of samples in the manifest.

__getstate__

def __getstate__() -> dict[str, Any]

Ship only the path (not the loaded rows) to DataLoader workers when parquet-backed.

from_parquet classmethod

def from_parquet(path: str | PathLike[str]) -> Manifest

Load a manifest from a parquet file.

Rows are read lazily, per (rank, worker) shard, so a DataLoader worker never loads the whole manifest into RAM.

generate classmethod

def generate(
    source: DataSource,
    index: str,
    fields: dict[str, Field],
    *,
    timeline_sampling: FixedRateSampling | None = None,
    fetch_block_size: int = 128,
    num_ranks: int = 1,
    num_workers_per_rank: int = 1,
    required_fields: list[str] | None = None,
    scan_max_workers: int | None = None,
) -> Manifest

Generate an unshuffled manifest for one epoch by scanning the source.

Scans the source, drops invalid samples, and unrolls one epoch in natural order (segment by segment, along the timeline) into fetch_group / emit_rank columns. To shuffle, call shuffle on the result — it re-orders cheaply without re-scanning. Persist with write_parquet.

See RerunIterableDataset for source / index / fields / timeline_sampling. fetch_block_size is the co-fetch / co-decode block size and num_ranks / num_workers_per_rank freeze the (rank, worker) assignment. required_fields are the field keys that must resolve to real data for a sample to be kept. scan_max_workers caps the concurrent scan queries against the server (defaults to 8); raise it to speed up scanning a large dataset.

shuffle

def shuffle(
    strategy: ShuffleStrategy | None = None,
    *,
    seed: int = 0,
) -> Manifest

Shuffle a scanned manifest into a new epoch order, without re-scanning the source.

Keeps the manifest's validated sample set and per-field decode ranges (the expensive scan result) and only recomputes the fetch_group / emit_rank schedule under a new strategy and seed. fetch_block_size and the (num_ranks, num_workers_per_rank) topology are inherited from this manifest. Returns a new manifest; this one is unchanged.

Scan once with generate, then call this per epoch (bumping seed) to get fresh orders for free.

PARAMETER	DESCRIPTION
strategy	The ShuffleStrategy to apply, e.g. BlockShuffle (the default), SampleShuffle, or NoShuffle. It fixes both the fetch order and the emission buffer, so passing the same strategy object here and to RerunIterableDataset guarantees a replay matches the live run. TYPE: ShuffleStrategy | None DEFAULT: None
seed	Seed for the block shuffle and the emission buffer. TYPE: int DEFAULT: 0

to_arrow

def to_arrow() -> Table

The manifest as an Arrow table (loading it from the backing parquet file if needed).

validate_topology

def validate_topology(
    num_ranks: int, num_workers_per_rank: int
) -> None

Raise if the run's topology differs from the one the manifest was frozen for.

worker_assignments

def worker_assignments(rank: int, worker: int) -> Table

Rows assigned to one (rank, worker), in fetch order.

For a parquet-backed manifest this reads only this shard from the file (predicate pushdown), not the whole manifest.

worker_plan

def worker_plan(
    rank: int, worker: int
) -> tuple[list[Table], ndarray]

Return this (rank, worker)'s fetch_groups (co-fetch / co-decode units, in fetch order) and emission pull order.

Reading a shard can hit disk, so the fetch groups and the emission order are both derived from a single read.

write_parquet

def write_parquet(
    path: str | PathLike[str],
    *,
    row_group_size: int = 1 << 16,
) -> None

Write the manifest to a zstd-compressed parquet file; the header rides in the schema metadata.

NoShuffle dataclass

Bases: ShuffleStrategy

Natural order (segment by segment, along the timeline): maximal fetch locality, no randomness.

emission_buffer

def emission_buffer() -> ShuffleBuffer | None

The post-decode buffer samples are emitted through, or None to emit in fetch order.

Only BlockShuffle defines one, because it is the only strategy whose fetch order stays deliberately correlated: emission is where its batches get decorrelated. SampleShuffle already fetches a uniform permutation, so a buffer would add nothing, and NoShuffle is a deterministic baseline that a buffer would only contaminate.

Owning the buffer here is what keeps a live run and a manifest replay in step: both read it off the same strategy object, so they cannot be configured with different buffers by accident.

NumericDecoder

Bases: ColumnDecoder[Tensor]

Decode Arrow numeric / list-of-numeric columns to tensors, one vectorized gather per batch.

Segment-blind: every request already carries its rows, so the gather runs across the whole fetch block at once rather than once per sample.

Windowed numeric lists require every resolved row to have the same width and return [T, D] (including [T, 1] for scalar components); a window with varying widths returns None. Unwindowed variable-width fields return one tensor per sample and require a padding or ragged-data collator for batching.

fill_latest_at property

fill_latest_at: bool

Whether this column's prefetch read latest-at-fills empty grid slots.

True for stateless columns (images, scalars): each grid slot wants the most recent value snapped from the real rows. Keyframe-aware fields use exact range queries instead, regardless of this value. The read is partitioned by this flag so it stays a global query argument per group rather than a per-column one.

prior_keyframe_path

def prior_keyframe_path(field_path: str) -> str | None

Sibling column whose non-null rows mark a re-entrant keyframe, or None.

Override on decoders that need the prefetch window anchored at the prior keyframe (compressed video). Default returns None.

RerunIterableDataset

Bases: IterableDataset[DecodedSample]

Iterable dataset backed by a catalog server.

Fetches fetch_block_size samples per server query and yields individual samples, so per-query overhead is amortized across many samples while the DataLoader controls the training batch size independently.

The index list is partitioned across DDP ranks and DataLoader workers internally. With shuffling enabled (default), the sample order is permuted once per epoch before partitioning; call set_epoch to re-seed between epochs. When a finite live dataset is consumed to exhaustion under DDP, wrap the training loop in DistributedDataParallel.join(): rank shards can have different lengths, especially when incomplete samples are skipped after sharding.

PARAMETER	DESCRIPTION
source	The dataset to read from (with optional segment filter). TYPE: DataSource
index	Timeline to iterate (e.g. "frame_nr"). TYPE: str
fields	Sample fields, keyed by output name. TYPE: dict[str, Field]
timeline_sampling	Required when index is a timestamp timeline; ignored for integer indices. Pass FixedRateSampling to sample on a fixed grid (e.g. 30 Hz). TYPE: FixedRateSampling | None DEFAULT: None
fetch_block_size	Number of samples to fetch per server query. Larger values amortize network overhead but use more memory. Defaults to 128. TYPE: int DEFAULT: 128
shuffle_strategy	The ShuffleStrategy that determines the order samples are fetched in, and — for BlockShuffle — the optional post-decode buffer samples are emitted through. Defaults to SampleShuffle; pass NoShuffle for natural order. TYPE: ShuffleStrategy | None DEFAULT: None
decode_threads	Fields to decode concurrently within each DataLoader worker. TYPE: int | None DEFAULT: None
max_consecutive_skipped_samples	Maximum number of consecutive incomplete samples to skip in each live iterator, independently for every rank and DataLoader worker. A valid sample resets the count. The next missing sample raises with total and per-field counts. Defaults to 1000; pass None to apply no limit. Manifest replay remains strict regardless of this setting. TYPE: int | None DEFAULT: _DEFAULT_MAX_CONSECUTIVE_SKIPPED_SAMPLES

sample_index property

sample_index: SampleIndex

The underlying SampleIndex.

__iter__

def __iter__() -> Iterator[DecodedSample]

Yield this worker's samples: replayed from a manifest, or fetched live from the catalog.

__len__

def __len__() -> int

Total number of samples across all segments.

from_manifest classmethod

def from_manifest(
    manifest: Manifest,
    source: DataSource,
    fields: dict[str, Field],
    *,
    decode_threads: int | None = None,
) -> RerunIterableDataset

Build a dataset that replays a frozen Manifest's sampling order.

The order, shards, and decode ranges all come from manifest; a manifest records only the field specs, not the objects, so the live connection and decoders are supplied here.

PARAMETER	DESCRIPTION
manifest	The frozen manifest to replay. Provides the sampling order, shards, and decode ranges. TYPE: Manifest
source	The live catalog connection. TYPE: DataSource
fields	The decoders, keyed by field name (a manifest records only their spec, not the objects). TYPE: dict[str, Field]
decode_threads	Fields to decode concurrently within each DataLoader worker. TYPE: int | None DEFAULT: None

RETURNS	DESCRIPTION
RerunIterableDataset	A dataset that yields samples in the manifest's recorded order.

set_epoch

def set_epoch(epoch: int) -> None

Set the epoch for shuffling (like DistributedSampler.set_epoch).

RerunMapDataset

Bases: Dataset[DecodedSample]

Map-style dataset backed by a catalog server.

Supports random access by global index, so it works with PyTorch's sampler ecosystem (DistributedSampler, WeightedRandomSampler, SubsetRandomSampler, ...). Shuffling and cross-worker partitioning are driven by the DataLoader's sampler.

For streaming iteration with internal shuffling, use RerunIterableDataset instead.

PARAMETER	DESCRIPTION
source	The dataset to read from (with optional segment filter). TYPE: DataSource
index	Timeline column to use as the sample index (e.g. "frame_nr"). TYPE: str
fields	Sample fields, keyed by output name. TYPE: dict[str, Field]
timeline_sampling	Required when index is a timestamp timeline; ignored for integer indices. Pass FixedRateSampling to sample on a fixed grid (e.g. 30 Hz). TYPE: FixedRateSampling | None DEFAULT: None
decode_threads	Fields to decode concurrently within each DataLoader worker. TYPE: int | None DEFAULT: None

Examples:

dataset = RerunMapDataset(
    source,
    "frame_nr",
    {"image": Field("/camera:Image:blob", decode=ImageDecoder())},
)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=8, sampler=sampler, num_workers=4)
for batch in loader:
    ...

sample_index property

sample_index: SampleIndex

The underlying SampleIndex.

__getitem__

def __getitem__(idx: int) -> DecodedSample

Fetch a single sample by global index (one server query).

__getitems__

def __getitems__(indices: list[int]) -> list[DecodedSample]

Fetch multiple samples by global index in a single server query.

PyTorch's DataLoader calls this automatically when present, so each batch round-trips once.

__len__

def __len__() -> int

Total number of samples across all segments.

from_manifest classmethod

def from_manifest(
    manifest: Manifest,
    source: DataSource,
    fields: dict[str, Field],
    *,
    decode_threads: int | None = None,
) -> RerunMapDataset

Build a map-style dataset over a frozen Manifest's validated samples.

This is a performance optimization only: it reuses the manifest's validated sample set and frozen decode ranges, so there is no live scan and no per-batch keyframe lookup.

Warning

The manifest's recorded order is not respected here. Ordering and cross-worker sharding stay with the DataLoader's sampler, as for any map-style dataset, so this cannot reproduce a manifest's run. For reproducible, resumable training use RerunIterableDataset.from_manifest.

PARAMETER	DESCRIPTION
manifest	The frozen manifest to read. Provides the validated sample set and decode ranges. TYPE: Manifest
source	The live catalog connection. TYPE: DataSource
fields	The decoders, keyed by field name (a manifest records only their spec, not the objects). TYPE: dict[str, Field]
decode_threads	Fields to decode concurrently within each DataLoader worker. TYPE: int | None DEFAULT: None

RETURNS	DESCRIPTION
RerunMapDataset	A dataset over the manifest's validated samples.

SampleIndex

Pre-computed description of the complete sample space.

Maps every segment's positional indices to concrete index values, accounting for the timeline strategy (integer or fixed-rate grid). Small enough to hold in memory for any realistic dataset.

PARAMETER	DESCRIPTION
segments	Per-segment metadata (window-adjusted index range + sample count). TYPE: list[SegmentMetadata]
ns_per_sample	For FixedRateSampling: nanoseconds between grid points. None for integer indices. TYPE: int | None DEFAULT: None
ns_dtype	Numpy dtype string used when materializing index values: "datetime64[ns]" for timestamp timelines, "timedelta64[ns]" for duration timelines, or None for plain integer indices. TYPE: str | None DEFAULT: None

is_duration property

is_duration: bool

Whether the index is a duration timeline.

is_timestamp property

is_timestamp: bool

Whether the index is a timestamp timeline.

ns_dtype property

ns_dtype: str | None

Numpy dtype for materialized index values, or None for integer indices.

ns_per_sample property

ns_per_sample: int | None

Nanoseconds between grid points for fixed-rate sampling, or None.

segment_offsets property

segment_offsets: ndarray

Cumulative offsets: segments[i] covers global indices [segment_offsets[i], segment_offsets[i + 1]).

segments property

segments: list[SegmentMetadata]

Per-segment metadata list.

total_samples property

total_samples: int

Total number of samples across all segments.

build staticmethod

def build(
    source: DataSource,
    index: str,
    fields: dict[str, Field],
    *,
    timeline_sampling: FixedRateSampling | None = None,
) -> SampleIndex

Build a SampleIndex from lightweight metadata queries.

PARAMETER	DESCRIPTION
source	Data source to build from. TYPE: DataSource
index	Name of the index timeline column. TYPE: str
fields	Field definitions, used for window-trim calculation. TYPE: dict[str, Field]
timeline_sampling	Required for timestamp and duration indices; ignored for integer indices. Pass FixedRateSampling for a regular grid. TYPE: FixedRateSampling | None DEFAULT: None

global_to_local

def global_to_local(
    idx: int,
) -> tuple[SegmentMetadata, IndexValue]

Map a global index [0, total_samples) to (segment, concrete_idx_value).

The returned index value is a plain int for integer timelines, a datetime64[ns] for timestamp timelines, and a timedelta64[ns] for duration timelines.

indices_in_range

def indices_in_range(lo: int, hi: int) -> Iterable[int]

Enumerate valid index values in [lo, hi].

For fixed-rate timelines the returned values walk down from hi in ns_per_sample steps (so they remain on the grid as long as hi is). For integer timelines, every value in [lo, hi] is returned. Values are plain int (ns-since-epoch for timestamp indices, ns count for duration indices); the caller casts the aggregated set to the right numpy dtype.

offset_index

def offset_index(
    index_value: IndexValue, offset: int | float
) -> IndexValue

Add an explicit field-window offset to an index value.

output_index_values

def output_index_values(
    index_value: IndexValue, field: Field
) -> tuple[IndexValue, ...]

Concrete index values requested by a field, preserving its explicit offset order.

resolve_local_index

def resolve_local_index(
    seg: SegmentMetadata, pos: int
) -> IndexValue

Convert a positional index within seg to a concrete index value.

pos is in [0, seg.num_samples). Returns datetime64[ns] for timestamp timelines, timedelta64[ns] for duration timelines, and a plain int for integer indices.

SampleShuffle dataclass

Bases: ShuffleStrategy

Uniform per-sample shuffle: maximal decorrelation, minimal fetch locality.

Every fetch scatters across all segments; prefer BlockShuffle when fetch throughput is the bottleneck.

emission_buffer

def emission_buffer() -> ShuffleBuffer | None

The post-decode buffer samples are emitted through, or None to emit in fetch order.

Only BlockShuffle defines one, because it is the only strategy whose fetch order stays deliberately correlated: emission is where its batches get decorrelated. SampleShuffle already fetches a uniform permutation, so a buffer would add nothing, and NoShuffle is a deterministic baseline that a buffer would only contaminate.

Owning the buffer here is what keeps a live run and a manifest replay in step: both read it off the same strategy object, so they cannot be configured with different buffers by accident.

SegmentMetadata dataclass

Per-segment metadata for sampling.

ShuffleStrategy

Bases: ABC

Determines the order in which an epoch's samples are fetched.

See the training guide for the trade-offs.

emission_buffer

def emission_buffer() -> ShuffleBuffer | None

The post-decode buffer samples are emitted through, or None to emit in fetch order.

Only BlockShuffle defines one, because it is the only strategy whose fetch order stays deliberately correlated: emission is where its batches get decorrelated. SampleShuffle already fetches a uniform permutation, so a buffer would add nothing, and NoShuffle is a deterministic baseline that a buffer would only contaminate.

Owning the buffer here is what keeps a live run and a manifest replay in step: both read it off the same strategy object, so they cannot be configured with different buffers by accident.

epoch_order abstractmethod

def epoch_order(
    sample_index: SampleIndex,
    *,
    fetch_block_size: int,
    seed: int,
) -> tuple[ndarray, ndarray]

Return (indices, block_bounds) for one epoch: every global sample index once, in emission order.

block_bounds are cumulative end positions of blocks within indices; each block must be a contiguous, segment-local span of the global index space.

VideoFrameDecoder

Bases: ColumnDecoder[DecodedValue], Generic[_OutputFormatT]

Compressed video random access via keyframe-aware fetching.

Anchors the decode window at the prior keyframe by consulting the sibling is_keyframe component on the VideoStream archetype, derived from Field.path (e.g. /cam:VideoStream:sample pairs with /cam:VideoStream:is_keyframe). The marker is populated by the user or by LazyChunkStream.collect(optimize=…), and lives in dedicated chunks separate from the video sample, so the lookup is cheap.

The sibling is_keyframe column is required. This makes every decode range deterministic rather than relying on an estimated GOP length.

Samples may be raw H.264 AVC1/AVCC (length-prefixed NAL units) or Annex B; the format is detected automatically per sample.

A batch's requests are grouped by GOP: each GOP's packets are extracted and fed through the codec once, and every requested frame is captured as it is emitted. A batch (or a later batch) whose window extends an earlier one reuses the open codec context and decodes only the new packets.

With output_format="rgb", a Field.window returns one frame per explicit offset as a [T, 3, H, W] tensor. With output_format="yuv420p", it returns a Yuv420Frame whose Y and UV planes remain in YUV form until collation and device transfer.

Returns None when a request's resolved window contains no decodable keyframe: the target precedes the entity's first frame in a multi-modal segment, or the first row was user-logged is_keyframe=true on a sample that isn't actually a codec keyframe (run optimize with fix_keyframe=True to re-derive markers from the encoded samples). Consumers must filter these out in their collate function before stacking.

fill_latest_at property

fill_latest_at: bool

Whether this column's prefetch read latest-at-fills empty grid slots.

True for stateless columns (images, scalars): each grid slot wants the most recent value snapped from the real rows. Keyframe-aware fields use exact range queries instead, regardless of this value. The read is partitioned by this flag so it stays a global query argument per group rather than a per-column one.

__getstate__

def __getstate__() -> dict[str, Any]

Drop the sessions: open codec contexts cannot be pickled. Cache stats restart per process.

__init__

def __init__(
    *,
    codec: str = "h264",
    max_decoder_sessions: int = 8,
    thread_count: int = 1,
    window_storage: Literal["copy", "view"] = "copy",
    output_format: Literal["rgb"] = "rgb",
) -> None

def __init__(
    *,
    codec: str = "h264",
    max_decoder_sessions: int = 8,
    thread_count: int = 1,
    window_storage: Literal["copy", "view"] = "copy",
    output_format: Literal["yuv420p"],
) -> None

def __init__(
    *,
    codec: str = "h264",
    max_decoder_sessions: int = 8,
    thread_count: int = 1,
    window_storage: Literal["copy", "view"] = "copy",
    output_format: Literal["rgb", "yuv420p"] = "rgb",
) -> None

Construct a decoder for a compressed video column.

PARAMETER	DESCRIPTION
codec	Video codec of the encoded samples (e.g. "h264"). TYPE: str DEFAULT: 'h264'
max_decoder_sessions	Upper bound on the number of live codec contexts kept in the LRU cache. Set to 0 to disable session reuse. TYPE: int DEFAULT: 8
thread_count	FFmpeg decode thread count. 0 leaves thread selection to FFmpeg, while 1 requests single-threaded decoding and remains the default for a predictable thread budget. For H.264 and H.265, values greater than 1 enable frame threading, which decodes several frames concurrently but prevents decoder-session reuse because delayed frames must be flushed at the end of each decode run. TYPE: int DEFAULT: 1
window_storage	How windowed outputs are materialized. "copy" returns independent tensors. "view" stores each unique decoded frame once per decode run and returns views for contiguous windows. Overlapping views share storage and must not be mutated in place before collation. TYPE: Literal['copy', 'view'] DEFAULT: 'copy'
output_format	"rgb" returns the existing [3, H, W] uint8 tensor. "yuv420p" returns a compact Yuv420Frame without CPU RGB conversion. Stack it with Yuv420Frame.stack in an application collate function, or use Yuv420Collator. Perform RGB conversion in the training process after transfer to the GPU. Combine it with window_storage="view" to decode overlapping windows directly into a shared YUV frame bank. TYPE: Literal['rgb', 'yuv420p'] DEFAULT: 'rgb'

decode

def decode(
    batch: FieldBatch, requests: Sequence[DecodeRequest]
) -> Sequence[Tensor | None]

def decode(
    batch: FieldBatch, requests: Sequence[DecodeRequest]
) -> Sequence[Yuv420Frame | None]

def decode(
    batch: FieldBatch, requests: Sequence[DecodeRequest]
) -> Sequence[DecodedValue | None]

Decode each request's frame or frame window, feeding every GOP once.

Yuv420Collator

Collate dataset mappings containing Yuv420Frame values.

Without device, YUV planes remain compact CPU uint8 tensors. This is the intended configuration with DataLoader(num_workers>0, pin_memory=True): convert each collated Yuv420Frame in the training process using Yuv420Frame.to_rgb.

Examples:

Custom collation can handle each non-video field however the application requires:

def collate(samples):
    return {
        "video": Yuv420Frame.stack([sample["video"] for sample in samples]),
        "state": custom_state_collation(samples),
    }

The convenience collator applies PyTorch's default collation to non-video fields:

loader = DataLoader(dataset, collate_fn=Yuv420Collator(), pin_memory=True)
for batch in loader:
    batch["video"] = batch["video"].to_rgb("cuda", dtype=torch.float16, non_blocking=True)

__call__

def __call__(
    samples: Sequence[Mapping[str, object]],
) -> dict[str, object]

Stack a batch of sample mappings, preserving YUV for device-side conversion.

Yuv420Frame dataclass

One frame, frame window, or collated batch in planar YUV420 form.

y has shape [..., 1, H, W] and uv has shape [..., 2, ceil(H / 2), ceil(W / 2)]. Both tensors are uint8. uv[..., 0, :, :] is U and uv[..., 1, :, :] is V. Leading dimensions may contain time and, after collation, batch.

Use Yuv420Frame.stack inside a custom collate function, or the convenience Yuv420Collator. The standard PyTorch collator does not know how to stack the planes and metadata.

clone

def clone() -> Yuv420Frame

Return an independent copy of both planes.

pin_memory

def pin_memory() -> Yuv420Frame

Pin both planes so a DataLoader can transfer them asynchronously.

stack classmethod

def stack(frames: Sequence[Yuv420Frame]) -> Yuv420Frame

Stack frames or windows along a new leading batch dimension.

This is the composable building block for application-specific collate functions. All inputs must have identical shapes and color metadata. Stacking materializes window views, so the returned batch no longer shares decoder frame-bank storage with its inputs.

to_rgb

def to_rgb(
    device: device | str | None = None,
    *,
    dtype: dtype = float32,
    normalize: bool = True,
    non_blocking: bool = False,
    color_space: Literal["bt601", "bt709", "bt2020"]
    | None = None,
    color_range: Literal["limited", "full"] | None = None,
) -> Tensor

Transfer and convert to a contiguous [..., 3, H, W] RGB tensor.

Conversion goes directly from CPU uint8 YUV into the requested floating-point dtype on device, so an intermediate RGB uint8 allocation is not required. If normalize is true, output values are in [0, 1]; otherwise they are in [0, 255].

Chroma is expanded with nearest-neighbor sampling to stay close to FFmpeg's default rgb24 conversion. Conversion can still differ by a few intensity levels because FFmpeg uses integer arithmetic.

color_space and color_range override the metadata carried by the decoded frame. Unspecified metadata falls back to BT.601 limited range with a warning, matching FFmpeg's conventional SD-video default.

tracing_scope

def tracing_scope(name: str) -> Iterator[None]

Open an OpenTelemetry span for the duration of a with block and propagate trace context into Rerun's Rust SDK.

Context-manager counterpart to with_tracing — use it to scope arbitrary blocks of code without extracting them into a function. Any Rust-side #[instrument] spans triggered from within will be parented under this span in Jaeger.

No-op unless TELEMETRY_ENABLED=true and an OTLP endpoint is configured (OTEL_EXPORTER_OTLP_TRACES_ENDPOINT or OTEL_EXPORTER_OTLP_ENDPOINT).

Examples:

for epoch in range(num_epochs):
    with tracing_scope(f"epoch {epoch}"):
        train_one_epoch(...)

with_tracing

def with_tracing(name: str) -> Callable[[F], F]

Wrap a function in an OpenTelemetry span and propagate trace context into Rerun's Rust SDK.

When enabled, creates a span named name, injects the W3C traceparent header into Rerun's shared ContextVar, and runs the wrapped function. Any Rust-side #[instrument] spans triggered from within (e.g. catalog queries) will be parented under this span in Jaeger.

For ad-hoc blocks that don't belong in a dedicated function, use tracing_scope instead.

No-op unless TELEMETRY_ENABLED=true and an OTLP endpoint is configured (OTEL_EXPORTER_OTLP_TRACES_ENDPOINT or OTEL_EXPORTER_OTLP_ENDPOINT).