hyrax.datasets.streaming_data_provider
======================================

.. py:module:: hyrax.datasets.streaming_data_provider

.. autoapi-nested-parse::

   A ``DataProvider``-like wrapper that encapsulates a streaming dataset.

   :class:`~hyrax.datasets.data_provider.DataProvider` is map-style — it assumes indexed
   random access (``resolve_data(idx)``, joins, caching, augmentation, splits), none of
   which a live stream supports. :class:`StreamingDataProvider` is its **iterable** sibling:
   it wraps a single streaming dataset (e.g.
   :class:`~hyrax.datasets.kafka_stream_dataset.KafkaStreamDataset`) and presents the same
   surface the rest of Hyrax expects — ``collate`` (shared via
   :class:`~hyrax.datasets.data_provider.CollationMixin`) and ``sample_data`` for model
   pre-flighting — 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
----------

.. autoapisummary::

   hyrax.datasets.streaming_data_provider.logger


Classes
-------

.. autoapisummary::

   hyrax.datasets.streaming_data_provider.StreamingDataProvider


Module Contents
---------------

.. py:data:: logger

.. py:class:: StreamingDataProvider(config: dict, request: dict)

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


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

   :param config: The Hyrax runtime configuration.
   :type config: dict
   :param request: 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).
   :type request: dict


   .. py:attribute:: config


   .. py:attribute:: data_request


   .. py:attribute:: friendly_name


   .. py:attribute:: primary_dataset


   .. py:attribute:: primary_dataset_id_field_name


   .. py:attribute:: primary_id_field


   .. py:attribute:: fields


   .. py:attribute:: _stream


   .. py:attribute:: prepped_datasets


   .. py:attribute:: custom_collate_functions
      :type:  dict


   .. py:attribute:: field_collate_functions
      :type:  dict


   .. py:method:: _register_field_collate_hooks()

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



   .. py:method:: _structure(sample: dict) -> dict

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



   .. py:method:: __iter__()

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



   .. py:method:: sample_data() -> dict

      Return one structured sample for model pre-flighting (``setup_model``).

      Peeks a single message from the stream without losing it (it is replayed in the
      first batch) and structures it like any other sample.



   .. py:method:: stop()

      Stop the underlying stream's iteration.