hyrax.datasets.lsdb_stream_dataset

Streaming dataset that reads rows from a HATS catalog via LSDB.

LSDBStreamDataset is a torch.utils.data.IterableDataset that wraps lsdb.streams.CatalogStream (a single finite pass over every partition) or lsdb.streams.InfiniteStream (endless random resampling of partitions), so a HATS catalog can be trained on or inferred over without ever materializing it in memory. This is the streaming counterpart to HyraxHATSDataset, which calls catalog.compute() and holds the whole table.

Rule of thumb: use stream_type = "catalog" for inference (every object is visited exactly once and the run ends on its own when the catalog is exhausted) and stream_type = "infinite" for training (batches keep arriving until you stop the session, so training is not bounded by a single pass).

Fixed-size batching. LSDB yields one pandas.DataFrame per chunk of partitions, so chunk sizes are ragged - a catalog may hand back 9 rows, then 10, then 171. This dataset buffers rows across chunks and splits large chunks, so every yielded batch holds exactly [data_loader] batch_size rows. Only the final batch, emitted when a finite stream is exhausted or stop() is called, may be short.

Specifying the catalog. data_location is an lsdb://<name> handle referring to an in-memory catalog registered with register_catalog().

import hyrax
import lsdb
from hyrax.datasets import LSDBStreamDataset

gaia = lsdb.open_catalog("https://data.lsdb.io/hats/gaia_dr3")
data_location = LSDBStreamDataset.register_catalog(
    "gaia_bright", gaia.query("phot_g_mean_mag < 20")
)

hy = hyrax.Hyrax()
hy.config["data_request"] = {
    "infer_stream": {
        "data": {
            "dataset_class": "LSDBStreamDataset",
            "data_location": data_location,
            "primary_id_field": "gaia_id",
            "fields": ["ra", "dec", "phot_g_mean_mag"],
        }
    }
}

with hy.infer_stream() as session:  # builds the provider + loader internally
    for batch, results in session:
        ...

Note

Nested columns are collated by default. Variable-length nested arrays are padded to the longest row in each batch, in the column’s own dtype, and accompanied by a boolean <field>_mask; fixed-length nested columns are stacked directly. Padding is whatever numpy zero-initializes that dtype to (0, or "" for a string column), so the mask is the only reliable record of what is real. Define a collate_<field> method on a subclass when a different representation is needed.

Warning

The stream owns a single in-process iterator, so the loader must run with num_workers = 0 (the default applied by dist_data_loader). A plain IterableDataset does no worker sharding, so with multiple workers each would build its own stream and every row would be emitted once per worker.

Keeping the consumer fed. Because the loader runs with num_workers = 0, PyTorch’s own prefetch_factor does nothing here and this dataset has to hide fetch latency itself. Three settings do that, and they stack:

use_dask_client

lsdb’s CatalogIterator submits the next chunk’s future before returning the current one, but only a real dask.distributed client makes that submission asynchronous - without one it computes inline, so every fetch blocks. Hyrax attaches to whatever client is already active, so create one before the run starts. On a single machine prefer Client(processes=False): it still gives real futures, but keeps results in-process rather than serializing every chunk back from a worker, which is a real cost for the nested frames light-curve catalogs produce.

partitions_per_chunk

How much work one fetch is worth. Raising it amortizes per-chunk overhead over more rows and, with a client, spreads a single chunk across more workers; lsdb recommends at least twice the worker count. The cost is memory.

Enable DEBUG logging on this module to see where the time actually goes - fetch, decode, and how long the consumer blocked. When prefetching is working, the consumer wait falls toward zero while the fetch time does not.

Attributes

logger
LSDB_URI_PREFIX
_CATALOG_REGISTRY

Classes

LSDBStreamDataset

Streams rows from a HATS catalog and yields fixed-size batches.

Module Contents

logger

LSDB_URI_PREFIX = 'lsdb://'

_CATALOG_REGISTRY: dict[str, object]

class LSDBStreamDataset(config: dict, data_location=None)

Bases: hyrax.datasets.dataset_registry.HyraxDataset, torch.utils.data.IterableDataset

Streams rows from a HATS catalog and yields fixed-size batches.

The stream is configured in [data_set.LSDBStreamDataset]; the catalog itself comes from the data_location of the data request, as an lsdb://<name> handle of the in-memory catalog registry.

Each row becomes a flat dict of column name to value (e.g. {"object_id": 2787..., "ra": 95.4, "dec": -36.3}); the wrapping StreamingDataProvider turns each flat sample into the structured form the collation + model machinery expect.

__init__()

Overall initialization for all Datasets which saves the config

Subclasses of HyraxDataset ought call this at the end of their __init__ like:

from hyrax.datasets import HyraxDataset

class MyDataset(HyraxDataset):
    def __init__(config):
        <your code>
        super().__init__(config)

If per tensor metadata is available, it is recommended that dataset authors create an astropy Table of that data, in the same order as their data and pass that metadata_table as shown below:

from hyrax.datasets import HyraxDataset
from astropy.table import Table

class MyDataset(HyraxDataset):
    def __init__(config):
        <your code>
        metadata_table = Table(<Your catalog data goes here>)
        super().__init__(config, metadata_table)

Parameters:

• config (dict, Optional) – The runtime configuration for hyrax

• metadata_table (Optional[Table], optional) – An Astropy Table with 1. the metadata columns desired for visualization AND 2. in the order your data will be enumerated.

• object_id_column_name (Optional[str], optional) – The name of the column containing object IDs. If None, uses the default from config or creates one from the ids() method.

data_location = ''

stream_type = ''

partitions_per_chunk

shuffle

seed = None

use_dask_client

dask_client_address = None

batch_size

_catalog

_stop

_iterator = None

_exhausted = False

_buffered: list[dict] = []

_owned_client = None

_register_getters(columns) → None

_register_nested_getters(nested_column, nested_subcolumns) → None

_register_nested_collators(nested_column, nested_subcolumns) → None

classmethod register_catalog(name: str, catalog) → str

Register an in-memory LSDB catalog and return its lsdb:// data_location.

Use this for catalogs that cannot be expressed as a path, such as the result of a crossmatch or query. The returned handle is an ordinary string, so it can be stored in the config and serialized with the rest of the runtime config.

Parameters:

• name (str) – Name to register the catalog under. Must not contain whitespace or /.

• catalog (lsdb.Catalog) – The catalog to stream from.

Returns:

The lsdb://<name> string to use as the request’s data_location.

Return type:

str

classmethod unregister_catalog(name: str) → None

Remove a registered catalog. Does nothing if the name is not registered.

classmethod clear_catalogs() → None

Remove every registered catalog.

classmethod registered_catalogs() → list[str]

Return the sorted names of all registered catalogs.

classmethod _lookup_catalog(name: str)

Return the catalog registered under name, or raise a helpful KeyError.

_resolve_dask_client()

Return the dask client to hand lsdb, or None to compute synchronously.

This is the difference between a stream that overlaps with training and one that does not. CatalogIterator.__next__ submits the next chunk before returning the current one, but with no client lsdb’s submit_next_partitions computes inline and hands back an already-resolved stand-in, so every fetch blocks.

_make_stream()

Build the LSDB stream over the catalog.

Built lazily (not in __init__) so that constructing the dataset stays cheap, so a dask client created around the run is picked up, and so tests can substitute a stream of plain DataFrames.

_ensure_iterator()

Return the shared stream iterator, creating it on first use.

CatalogStream.__iter__ hands back a new iterator with a freshly spawned RNG on every call, so successive passes of __iter__() must share a single one or the rows one pass pushed back would come from a different traversal than the next.

stop()

Signal __iter__() to flush any pending rows and stop fetching chunks.

LSDB offers no cancellation or timeout, so a __iter__() already blocked inside a chunk computation cannot be interrupted; the stop takes effect at the next chunk boundary. In the usual flow this is invisible, because the session calls stop() from the same thread after iteration has ended. To end an infinite stream promptly, break out of the session loop rather than calling stop() from another thread. A smaller partitions_per_chunk bounds the delay.

close()

Stop iteration and close the dask client this dataset created, if any. Idempotent.

Only a client created here from dask_client_address is closed. A client attached via the active-client fallback (get_client()) is owned by whoever created it - closing it here would pull it out from under the rest of the session. Skipping this for an owned client leaves its scheduler registration and background threads running after the stream is done with it.

_owned_client is None is the guard, so a second call is a no-op rather than a double close.

__len__()

A stream has no length.

Defined only so HyraxDataset.__init_subclass__ (which requires a __len__ attribute) accepts the class. The iterable branch of dist_data_loader never calls it.

__iter__()

Yield list[dict] batches of exactly batch_size rows.

Rows are buffered across chunks and large chunks are split, so only the final batch - emitted when a finite stream is exhausted or stop() is set - may be shorter than batch_size.

Prefetching changes when chunks are computed, never which rows come out or in what order: the loop below still refuses to take a new chunk once stop() is set, and every prefetched row the consumer does not reach is pushed back.