Data¶

tinax.data provides a read/write interface, deterministic process-aware training and evaluation pipelines, dataset splitting, random-access sources, and explicit multiprocessing worker lifetime.

Grain multiprocessing requires caller-parsed Abseil flags. Tinax validates this requirement but never mutates global flags.

Both the ArrayRecord and Parquet pairs return the same random-access grain.MapDataset, so they compose identically with split_by_column, split_random, training_batches, and evaluation_batches. read_array_record/write_array_record need the optional array-record package installed separately; pyarrow for the Parquet pair is a core dependency.

read_array_record ¶

read_array_record(path: str | PathLike[str]) -> MapDataset

Load one ArrayRecord file's JSON-encoded rows into a lazy, randomly-accessible dataset.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Path to an existing ArrayRecord file.	required

Returns:

Type	Description
`MapDataset`	A `grain.MapDataset` decoding each row from JSON on access.

Raises:

Type	Description
`TypeError`	If `path` is not a string or `os.PathLike`.
`FileNotFoundError`	If no file exists at `path`.

write_array_record ¶

write_array_record(
    path: str | PathLike[str],
    records: Iterable[Mapping[str, Any]],
    *,
    group_size: int = 1,
) -> None

Serialize an iterable of JSON-encodable mapping rows into an ArrayRecord file.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination file path.	required
`records`	`Iterable[Mapping[str, Any]]`	Iterable of JSON-encodable mappings, one per row.	required
`group_size`	`int`	Number of records grouped per compressed chunk. Must be a positive integer.	`1`

Raises:

Type	Description
`TypeError`	If `path` is not a string or `os.PathLike`, or `group_size` is not an integer.
`ValueError`	If `group_size` is not positive.
`ImportError`	If `array-record` is not installed.

read_parquet_record ¶

read_parquet_record(
    path: str | PathLike[str],
) -> MapDataset

Load one Parquet file's rows into a lazy, randomly-accessible dataset.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Path to an existing Parquet file.	required

Returns:

Type	Description
`MapDataset`	A `grain.MapDataset` over the file's rows, each a `dict[str, Any]`.

Raises:

Type	Description
`TypeError`	If `path` is not a string or `os.PathLike`.
`FileNotFoundError`	If no file exists at `path`.

write_parquet_record ¶

write_parquet_record(
    path: str | PathLike[str],
    records: Sequence[Mapping[str, Any]],
) -> None

Serialize a sequence of mapping rows into a Parquet file.

Unlike write_array_record, records must be a finite, in-memory sequence: Parquet's columnar layout requires the full column set up front and cannot stream row by row.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination file path. Overwritten if it already exists.	required
`records`	`Sequence[Mapping[str, Any]]`	Non-empty sequence of mapping rows sharing one column schema.	required

Raises:

Type	Description
`TypeError`	If `path` is not a string or `os.PathLike`, or `records` is not a sequence of mappings.
`ValueError`	If `records` is empty.

split_by_column ¶

split_by_column(
    dataset: MapDataset[dict[str, Any]],
    column: str,
    *,
    values: tuple[str, str] = ("train", "test"),
    drop_column: bool = True,
) -> tuple[
    MapDataset[dict[str, Any]], MapDataset[dict[str, Any]]
]

Partition dataset by an exact two-value discriminator column.

Every row's column value must be one of values; a row holding anything else would otherwise be silently dropped by a naive filter-based split. The returned partitions are re-indexed sources with correct len and indexing, so they compose with the length-based training_batches and evaluation_batches pipelines.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[dict[str, Any]]`	Finite `grain.MapDataset` of mapping rows.	required
`column`	`str`	Name of the column carrying the discriminator value.	required
`values`	`tuple[str, str]`	The (first, second) partition values.	`('train', 'test')`
`drop_column`	`bool`	Whether to drop `column` from both returned partitions.	`True`

Returns:

Type	Description
`tuple[MapDataset[dict[str, Any]], MapDataset[dict[str, Any]]]`	The two partitions of `dataset`, in the order given by `values`.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `column` is not a string, or `drop_column` is not a bool.
`ValueError`	If `values` does not hold two distinct entries, or any row's `column` value is not one of `values`.

split_random ¶

split_random(
    dataset: MapDataset[dict[str, Any]],
    train_ratio: float,
    *,
    seed: int,
    drop_column: str | None = None,
) -> tuple[
    MapDataset[dict[str, Any]], MapDataset[dict[str, Any]]
]

Shuffle dataset and cut it into two partitions at train_ratio.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[dict[str, Any]]`	Finite `grain.MapDataset` to split.	required
`train_ratio`	`float`	Fraction of records in the first partition; must be strictly between 0 and 1 so neither partition is degenerately empty.	required
`seed`	`int`	Integer seed for the shuffle.	required
`drop_column`	`str \| None`	Optional column name to drop from both partitions afterward.	`None`

Returns:

Type	Description
`tuple[MapDataset[dict[str, Any]], MapDataset[dict[str, Any]]]`	The (first, second) partitions.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `train_ratio` is not a float, `seed` is not an integer, or `drop_column` is not a string or `None`.
`ValueError`	If `train_ratio` is not strictly between 0 and 1, or the dataset holds fewer than two records.

InMemorySource ¶

InMemorySource(records: Iterable[T], *, source_id: str)

Snapshot an outer iterable while preserving references to its possibly mutable records.

A deterministic random-access source supporting len() and integer indexing. The outer iterable is copied into a tuple at construction, but individual records are held by reference and are not deep copied.

Parameters:

Name	Type	Description	Default
`records`	`Iterable[T]`	Iterable of records to snapshot into the source.	required
`source_id`	`str`	Stable, non-empty identifier used in source representations.	required

Raises:

Type	Description
`TypeError`	If `source_id` is not a string.
`ValueError`	If `source_id` is empty.

source_id `property` ¶

source_id: str

Return the caller-assigned stable identity used in source representations.

training_batches ¶

training_batches(
    dataset: MapDataset[T],
    *,
    shard_options: ShardOptions,
    seed: int,
    shuffle: bool,
    num_epochs: int | None,
    batch_size: int,
    batch_fn: BatchTransform[T, BatchT],
) -> MapDataset[BatchT]

Build deterministic full local batches with equal process step counts in every training epoch.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[T]`	Finite `grain.MapDataset` describing one epoch of records.	required
`shard_options`	`ShardOptions`	Per-process shard selection. `drop_remainder` must be `True` so every process yields the same number of steps.	required
`seed`	`int`	Integer seed for deterministic ordering and shuffling.	required
`shuffle`	`bool`	Whether to shuffle records within each epoch.	required
`num_epochs`	`int \| None`	Number of epochs to repeat, or `None` for an unbounded stream.	required
`batch_size`	`int`	Local (per-process) batch size. Must be positive.	required
`batch_fn`	`BatchTransform[T, BatchT]`	Callable collating a sequence of records into one batch.	required

Returns:

Type	Description
`MapDataset[BatchT]`	A `grain.MapDataset` of collated batches, sharded to this process and
`MapDataset[BatchT]`	repeated for `num_epochs` with per-epoch reseeding.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `shard_options` is not a `grain.sharding.ShardOptions`, `seed` is not an integer, `shuffle` is not a bool, `batch_fn` is not callable, or a size is a boolean.
`ValueError`	If `shard_options.drop_remainder` is not `True`, a size is not positive, the dataset is not one finite epoch, or it holds fewer than one complete global batch.

evaluation_batches ¶

evaluation_batches(
    dataset: MapDataset[T],
    *,
    shard_options: ShardOptions,
    batch_size: int,
    batch_fn: BatchTransform[T, BatchT],
) -> MapDataset[BatchT]

Build one ordered finite pass that keeps remainder records and may have unequal process step counts.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[T]`	Finite `grain.MapDataset` describing the evaluation records.	required
`shard_options`	`ShardOptions`	Per-process shard selection. `drop_remainder` must be `False` so remainder records are kept.	required
`batch_size`	`int`	Local (per-process) batch size. Must be positive.	required
`batch_fn`	`BatchTransform[T, BatchT]`	Callable collating a sequence of records into one batch.	required

Returns:

Type	Description
`MapDataset[BatchT]`	A `grain.MapDataset` of collated batches for this process's shard, in order
`MapDataset[BatchT]`	and without dropping the final partial batch.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `shard_options` is not a `grain.sharding.ShardOptions`, `batch_fn` is not callable, or `batch_size` is a boolean.
`ValueError`	If `shard_options.drop_remainder` is not `False`, `batch_size` is not positive, or the dataset is not one finite epoch.

open_multiprocessing_iterator ¶

open_multiprocessing_iterator(
    dataset: IterDataset[T],
    *,
    num_workers: int,
    per_worker_buffer_size: int = 1,
    worker_init_fn: Callable[[int, int], None]
    | None = None,
    sequential_slice: bool = False,
) -> Iterator[DatasetIterator[T]]

Open a Grain multiprocessing iterator and close every worker when the context exits.

A context manager yielding a dataset iterator backed by worker processes; on exit the iterator and all workers are closed.

Parameters:

Name	Type	Description	Default
`dataset`	`IterDataset[T]`	`grain.IterDataset` to iterate with prefetching workers.	required
`num_workers`	`int`	Number of worker processes. `0` runs in the main process.	required
`per_worker_buffer_size`	`int`	Prefetch buffer size per worker. Must be positive.	`1`
`worker_init_fn`	`Callable[[int, int], None] \| None`	Optional callable run in each worker as `worker_init_fn(worker_index, worker_count)`.	`None`
`sequential_slice`	`bool`	Whether workers slice the dataset sequentially.	`False`

Yields:

Type	Description
`DatasetIterator[T]`	A `grain.DatasetIterator` over the prefetched dataset.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.IterDataset`, an integer argument is a boolean, `worker_init_fn` is not callable or `None`, or `sequential_slice` is not a bool.
`ValueError`	If `num_workers` is negative or `per_worker_buffer_size` is not positive.
`RuntimeError`	If `num_workers` is nonzero but Abseil flags are not parsed; enter through `absl.app.run`.

Data¶

read_array_record ¶

write_array_record ¶

read_parquet_record ¶

write_parquet_record ¶

split_by_column ¶

split_random ¶

InMemorySource ¶

source_id property ¶

training_batches ¶

evaluation_batches ¶

open_multiprocessing_iterator ¶

source_id `property` ¶