grain.sources module

APIs for reading data from various file formats.

List of Members

class grain.sources.RandomAccessDataSource(*args, **kwargs)

Interface for datasets where storage supports efficient random access.

This Protocol defines the contract for any custom data source injected into the PyGrain pipeline. Implementations do not need to inherit from this class directly; they only need to implement the required structural methods (__len__ and __getitem__).

Notes: Checkpointing: If used with DataLoader, __repr__ has to be additionally implemented to support checkpointing.

Multiprocessing: If used with multiprocessing, the instance must be fully picklable.

Example

Implementing a minimal, checkpoint-safe custom data source:

from grain.sources import RandomAccessDataSource

class MyInMemorySource:
  def __init__(self, data: list):
    self._data = data
  def __len__(self) -> int:
    return len(self._data)
  def __getitem__(self, index: int):
    return self._data[index]
  def __repr__(self) -> str:
    # Required for PyGrain checkpointing with DataLoader
    return f"MyInMemorySource(size={len(self)})"

source = MyInMemorySource(["a", "b", "c"])
# source satisfies the RandomAccessDataSource protocol.
assert isinstance(source, RandomAccessDataSource)

__getitem__(index)

Returns the value for the given index.

This method must be thread-safe and deterministic.

Note that a number of sources take SupportsIndex instead of int for index. Such sources will still support int index and pass the isinstance check with this protocol, but all new source implementations should use int directly.

Parameters:

index (int) – An integer in [0, len(self)-1].

Returns:

The corresponding record. File data sources often return the raw bytes but records can be any Python object.

Return type:

T

__len__()

Returns the total number of records in the data source.

Returns:

The total count of accessible records.

Return type:

int

class grain.sources.ArrayRecordDataSource(*args, **kwargs)

Data source for ArrayRecord files.

Reads serialized byte records from one or more ArrayRecord files or file instructions with random access indexing support.

Example

Writing to and reading records from an ArrayRecord file using MapDataset:

from array_record.python import array_record.python.array_record_data_source as array_record_module
import grain

# Write records to an ArrayRecord file.
writer = array_record_module.ArrayRecordWriter(
    "/tmp/data.array_record", "group_size:1"
)
writer.write(b"example_byte_record_1")
writer.write(b"example_byte_record_2")
writer.close()

# Load the ArrayRecord file using ArrayRecordDataSource.
source = grain.sources.ArrayRecordDataSource(
    ["/tmp/data.array_record"]
)
parent_ds = grain.MapDataset.source(source)
print(len(parent_ds))
# 2
print(parent_ds[0])
# b'example_byte_record_1'

Parameters:

• paths (array_record.python.array_record_data_source.PathLikeOrFileInstruction | Sequence[array_record.python.array_record_data_source.PathLikeOrFileInstruction])

• reader_options (dict[str, str] | None)

__init__(paths, reader_options=None)

Creates a new ArrayRecordDataSource object.

See array_record.ArrayRecordDataSource for more details.

Parameters:

• paths (array_record.python.array_record_data_source.PathLikeOrFileInstruction | Sequence[array_record.python.array_record_data_source.PathLikeOrFileInstruction]) – A single path/FileInstruction or list of paths/FileInstructions.

• reader_options (dict[str, str] | None) – a dict[str, str] to be passed when creating a reader. For example, {index_storage_option:”in_memory”} stores the reader indices in memory versus {index_storage_option:”offloaded”} stores the indices on disk to save memory usage.

Raises:

ValueError – If reader_options is provided but not supported by the underlying ArrayRecord reader version.

class grain.sources.SharedMemoryDataSource(elements=None, *, name=None)

Simple in-memory data source for sequences that is sharable among multiple processes.

Note

This constrains storable values to only the int, float, bool, str (less than 10M bytes each), bytes (less than 10M bytes each), and None built-in data types. It also notably differs from the built-in list type in that these lists can not change their overall length (i.e. no append, insert, etc.)

Example

Sharing an in-memory sequence across multiple worker processes without duplicating memory:

import grain

# Store the sequence in OS shared memory. Unlike a standard Python list
# that is copied into every worker process when pickled, only a
# lightweight shared memory reference is sent to each worker.
data = [10, 20, 30, 40]
source = grain.sources.SharedMemoryDataSource(data)

# Read and transform elements across 2 worker processes.
ds = (
    grain.MapDataset.source(source)
    .map(lambda x: x * 2)
    .to_iter_dataset()
)

print(list(ds))
# [20, 40, 60, 80]

# Clean up the shared memory block when done.
source.close()
source.unlink()

Parameters:

• elements (Sequence[Any] | None)

• name (str | None)

__init__(elements=None, *, name=None)

Creates a new InMemoryDataSource object.

Parameters:

• elements (Sequence[Any] | None) – The elements for the sharable list.

• name (str | None) – The name of the datasource.

Raises:

ValueError – If neither elements nor name is provided.

class grain.sources.RangeDataSource(start, stop, step)

Range data source, similar to python range() function.

Produces a sequence of integers from start to stop with a given step, supporting efficient indexing and length lookup without loading data into memory.

Example

Creating a range data source and accessing elements:

import grain

# Create a range data source
source = grain.sources.RangeDataSource(start=0, stop=10, step=2)

# Create a MapDataset from the source
ds = grain.MapDataset.source(source)

# Print the length of the dataset
print(len(ds))
# 5

# Print all elements in the dataset
print(list(ds))
# [0, 2, 4, 6, 8]

Parameters:

• start (int)

• stop (int)

• step (int)

__init__(start, stop, step)

Initializes the RangeDataSource.

Parameters:

• start (int) – The start value of the range sequence.

• stop (int) – The stop boundary of the range sequence (exclusive).

• step (int) – The step increment between consecutive elements. Must not be 0.

Raises:

AssertionError – If step is 0 or if computed length is negative.