nemo_curator.stages.interleaved.utils.materialization

Module Contents

Classes

Name	Description
`_ClassifiedRows`	-

Functions

Name	Description
`_build_global_range_index`	Resolve filesystem paths and build per-filesystem deduplicated range indices.
`_build_image_mask`	-
`_classify_rows`	Partition pending image rows into three I/O strategy groups.
`_extract_tiff_frame`	Extract a single frame from a multi-frame TIFF, returning it as a single-frame TIFF.
`_fill_direct_read_rows`	Read each direct file once, share bytes across all rows referencing it.
`_fill_materialized_bytes`	-
`_fill_range_read_rows`	Batch byte-range reads grouped by filesystem, calling cat_ranges once per backend.
`_fill_tar_extract_rows`	Open each tar once and extract all needed members sequentially.
`_get_frame_index`	-
`_init_materialization_buffers`	-
`_read_direct_file`	-
`_scatter_range_blobs`	Distribute deduplicated range-read results, extracting TIFF frames as needed.
`_task_with_dataframe`	-
`materialize_task_binary_content`	Return a task with image-row binary content materialized from source_ref.

Data

_SRC_PARSE_COLS

_TAR_EXTENSIONS

API

class nemo_curator.stages.interleaved.utils.materialization._ClassifiedRows()

Bases: NamedTuple

direct_read

dict[str, list[int]]

missing

list[int]

range_read

dict[str, list[tuple[int, str, int, int, int | None]]]

tar_extract

dict[str, list[tuple[int, str, int | None]]]

nemo_curator.stages.interleaved.utils.materialization._build_global_range_index(
    groups: dict[str, list[tuple[int, str, int, int, int | None]]],
    storage_options: dict[str, object],
    error_values: list[str | None]
) -> list[tuple[object, dict[tuple[str, int, int], list[tuple[int, str, int | None]]]]]

Resolve filesystem paths and build per-filesystem deduplicated range indices.

Returns a list of (fs, unique_ranges) pairs — one entry per distinct filesystem. unique_ranges maps (fs_path, offset, size) to the rows that need that byte range. Returns an empty list if no valid paths exist.

Paths are grouped by filesystem instance (fsspec caches instances per protocol and storage options), so cat_ranges is called once per connection pool. Batches that span multiple storage backends (e.g. local + S3) are handled correctly — each backend gets its own call.

nemo_curator.stages.interleaved.utils.materialization._build_image_mask(
    df: pandas.DataFrame,
    only_missing_binary: bool,
    image_content_types: tuple[str, ...] | None
) -> pandas.Series

nemo_curator.stages.interleaved.utils.materialization._classify_rows(
    df: pandas.DataFrame,
    image_mask: pandas.Series
) -> nemo_curator.stages.interleaved.utils.materialization._ClassifiedRows

Partition pending image rows into three I/O strategy groups.

tar_extract: has member name but no byte_offset (must open tar and extractfile)
range_read: has member + byte_offset + byte_size (can use fs.cat_ranges)
direct_read: no member (path is the file itself)
missing: path is None/NaN

nemo_curator.stages.interleaved.utils.materialization._extract_tiff_frame(
    tiff_bytes: bytes,
    frame_index: int
) -> bytes | None

Extract a single frame from a multi-frame TIFF, returning it as a single-frame TIFF.

Returns the raw bytes unchanged if the data is not a TIFF.

nemo_curator.stages.interleaved.utils.materialization._fill_direct_read_rows(
    groups: dict[str, list[int]],
    storage_options: dict[str, object],
    binary_values: list[object],
    error_values: list[str | None]
) -> None

Read each direct file once, share bytes across all rows referencing it.

nemo_curator.stages.interleaved.utils.materialization._fill_materialized_bytes(
    df: pandas.DataFrame,
    image_mask: pandas.Series,
    storage_options: dict[str, object],
    binary_values: list[object],
    error_values: list[str | None]
) -> None

nemo_curator.stages.interleaved.utils.materialization._fill_range_read_rows(
    groups: dict[str, list[tuple[int, str, int, int, int | None]]],
    storage_options: dict[str, object],
    binary_values: list[object],
    error_values: list[str | None]
) -> None

Batch byte-range reads grouped by filesystem, calling cat_ranges once per backend.

Deduplicates identical (path, offset, size) tuples globally so the same byte range is only fetched once even if referenced by multiple rows. Correctly handles batches that span multiple storage backends (e.g. local files mixed with remote files) by dispatching each group to its own fs.

nemo_curator.stages.interleaved.utils.materialization._fill_tar_extract_rows(
    groups: dict[str, list[tuple[int, str, int | None]]],
    storage_options: dict[str, object],
    binary_values: list[object],
    error_values: list[str | None]
) -> None

Open each tar once and extract all needed members sequentially.

nemo_curator.stages.interleaved.utils.materialization._get_frame_index(
    df: pandas.DataFrame,
    idx: int
) -> int | None

nemo_curator.stages.interleaved.utils.materialization._init_materialization_buffers(
    df: pandas.DataFrame
) -> tuple[list[object], list[str | None]]

nemo_curator.stages.interleaved.utils.materialization._read_direct_file(
    path: str,
    storage_options: dict[str, object]
) -> bytes | None

nemo_curator.stages.interleaved.utils.materialization._scatter_range_blobs(
    blobs: list[object],
    range_keys: list[tuple[str, int, int]],
    unique_ranges: dict[tuple[str, int, int], list[tuple[int, str, int | None]]],
    binary_values: list[object],
    error_values: list[str | None]
) -> None

Distribute deduplicated range-read results, extracting TIFF frames as needed.

nemo_curator.stages.interleaved.utils.materialization._task_with_dataframe(
    task: nemo_curator.tasks.InterleavedBatch,
    df: pandas.DataFrame
) -> nemo_curator.tasks.InterleavedBatch

nemo_curator.stages.interleaved.utils.materialization.materialize_task_binary_content(
    task: nemo_curator.tasks.InterleavedBatch,
    io_kwargs: dict[str, object] | None = None,
    only_missing_binary: bool = True,
    image_content_types: tuple[str, ...] | None = None
) -> nemo_curator.tasks.InterleavedBatch

Return a task with image-row binary content materialized from source_ref.

Dispatches to three I/O strategies based on source_ref contents:

range_read: byte_offset + byte_size present -> batched fs.cat_ranges()
tar_extract: member present, no byte range -> open tar + extractfile
direct_read: no member -> read file directly

nemo_curator.stages.interleaved.utils.materialization._SRC_PARSE_COLS = ('_src_path', '_src_member', '_src_byte_offset', '_src_byte_size', '_src_frame_i...

nemo_curator.stages.interleaved.utils.materialization._TAR_EXTENSIONS = ('.tar', '.tar.gz', '.tgz')