hyrax.datasets.streaming_data_provider

A DataProvider-like wrapper that encapsulates a streaming dataset.

DataProvider is map-style — it assumes indexed random access (resolve_data(idx), joins, caching, augmentation, splits), none of which a live stream supports. StreamingDataProvider is its iterable sibling: it wraps a single streaming dataset (e.g. KafkaStreamDataset) and presents the collate surface the rest of Hyrax expects (shared via CollationMixin) while delegating iteration to the wrapped stream.

The wrapped stream is a dumb decoder that yields list[dict] batches of flat sample dicts. This provider owns the structuring: it extracts the object id (via primary_id_field) and groups model inputs (via fields) under the request friendly-name, producing the per-sample shape collate expects:

{"object_id": str, "<friendly_name>": {field: np.ndarray, ...}}

The single source of truth for primary_id_field and fields is [data_request].

Attributes

logger

Classes

StreamingDataProvider

Wrap a single streaming IterableDataset behind a DataProvider-like surface.

Module Contents

logger

class StreamingDataProvider(config: dict, request: dict)

Bases: hyrax.datasets.data_provider.CollationMixin, torch.utils.data.IterableDataset

Wrap a single streaming IterableDataset behind a DataProvider-like surface.

Parameters:

• config (dict) – The Hyrax runtime configuration.

• request (dict) – A single-entry data request group (friendly_name -> definition). The definition must specify dataset_class (a registered IterableDataset) and primary_id_field; fields is optional (derived from the first sample when omitted).

config

data_request

dataset_getters

friendly_name

primary_dataset

primary_dataset_id_field_name

primary_id_field

fields

_stream

prepped_datasets

custom_collate_functions: dict

field_collate_functions: dict

_register_field_collate_hooks()

Detect collate_<field> methods on the wrapped stream for each field.

_structure(sample: dict) → dict

Turn a flat decoded sample into the per-sample shape collate expects.

__iter__()

Yield list[dict] batches of structured samples for collate_fn.

stop()

Stop the underlying stream’s iteration.

close()

Release the underlying stream’s resources (e.g. close the Kafka consumer).

Signalling stop() is not enough for a stream that holds a connection: a suspended __iter__ generator never resumes to run its own cleanup. Not every streaming dataset needs teardown, so this is a no-op when the stream defines none.