nemo_curator.stages.interleaved.utils.schema

Centralized schema utilities for interleaved IO readers and writers.

All arrow-based readers/writers share these functions for type reconciliation and schema alignment (null-fill + reorder).

Module Contents

Functions

Name	Description
`align_table`	Pad, reorder, and cast table to match target exactly.
`reconcile_schema`	Build a schema with canonical types for reserved columns and inferred types for passthrough.
`resolve_schema`	Return the effective schema from user-supplied schema or overrides.

Data

_LARGE_COMPAT

API

nemo_curator.stages.interleaved.utils.schema.align_table(
    table: pyarrow.Table,
    target: pyarrow.Schema
) -> pyarrow.Table

Pad, reorder, and cast table to match target exactly.

Columns in target absent from table are added as null arrays.
Columns in table absent from target are dropped.
Column order matches target.

Reserved INTERLEAVED_SCHEMA columns allow safe=False casts so that explicit large↔small type overrides work (e.g. large_string→string for Parquet compat). Passthrough (user-defined) columns always use safe=True so that overflow errors surface rather than silently corrupt data (e.g. large_string→string on a >2 GB column).

nemo_curator.stages.interleaved.utils.schema.reconcile_schema(
    inferred: pyarrow.Schema
) -> pyarrow.Schema

Build a schema with canonical types for reserved columns and inferred types for passthrough.

Avoids unsafe downcasts (e.g. large_string -> string) that cause offset overflow on large tables read via the pyarrow backend.

nemo_curator.stages.interleaved.utils.schema.resolve_schema(
    schema: pyarrow.Schema | None,
    overrides: dict[str, pyarrow.DataType] | None
) -> pyarrow.Schema | None

Return the effective schema from user-supplied schema or overrides.

Priority: schema > overrides merged on top of INTERLEAVED_SCHEMA > None.

If schema is provided and overrides is also provided, overrides are ignored and a warning is emitted. Returns None if both are None.

nemo_curator.stages.interleaved.utils.schema._LARGE_COMPAT: dict[tuple[DataType, DataType], DataType] = {(pa.large_string(), pa.string()): pa.large_string(), (pa.large_binary(), pa.bin...