Torch¶

def catalog_batch_stream_for_epoch(
    asset_ref: "CatalogAssetRef | CatalogAssetSnapshotRef | CatalogAsset",
    epoch: int,
    *,
    columns: Sequence[str] | None = None,
    catalog_batch_size: int = 4096,
    shuffle_enabled: bool = True,
    shard: CatalogReadShard | None = None,
) -> CatalogBatchStreamParams:
    """Build a per-epoch catalog batch stream spec (explicit epoch → new stream)."""
    return CatalogBatchStreamParams.from_asset_ref(
        asset_ref,
        columns=columns,
        catalog_batch_size=catalog_batch_size,
        shuffle={"enabled": shuffle_enabled, "seed": epoch},
        shard=shard,
    )

async def async_iter_catalog_torch_samples(
    catalog: "CatalogAPI",
    stream: CatalogBatchStreamParams,
    *,
    sample_format: SampleFormat = "dict",
    feature_columns: list[str] | None = None,
    target_column: str | None = None,
    transform: Callable | None = None,
    row_preprocess: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    limit: int | None = None,
    batch_size: int = 32,
    execution_options: TorchExecutionOptions | None = None,
    stats: BatchStreamStats | None = None,
) -> AsyncIterator[Any]:
    """Async iterator over torch samples from a catalog batch stream.

    Caller-owned ``CatalogAPI`` — no background thread and no adapter-spawned
    ``AsyncNexus``. Demonstrated in ``archive/examples/nexus/catalog_torch_train.py``
    (Layer 2 / power-user); prefer ``to_torch_dataloader`` (example 07) for training UX.

    Each full pass opens a new ``catalog.iter_batches`` stream (same per-epoch
    contract as the sync ``DataLoader`` bridge).

    With ``decode='vectorized'``, each yielded value is one micro-batch (same as
    ``to_torch_dataloader(..., batch_size=N)``).
    """
    stream_stats = stats if stats is not None else BatchStreamStats()
    stream_stats.stream_iterations += 1
    opts = resolve_execution_options(execution_options)
    use_vectorized = opts.decode == "vectorized" and transform is None and batch_size >= 1
    emitted = 0
    async for record_batch in catalog.iter_batches(**stream.iter_batches_kwargs()):
        stream_stats.record_batch(record_batch)
        for sample in _iter_samples_from_record_batch(
            record_batch,
            sample_format=sample_format,
            feature_columns=feature_columns,
            target_column=target_column,
            transform=transform,
            row_preprocess=row_preprocess,
            execution_options=opts,
            micro_batch_size=batch_size if use_vectorized else None,
        ):
            if limit is not None and emitted >= limit:
                stream_stats.consumer_stopped_early = True
                return
            stream_stats.torch_steps += 1
            emitted += 1
            yield sample

async def load_catalog_record_batches(
    catalog: "CatalogAPI",
    *,
    asset_ref: Any = None,
    columns: list[str] | None = None,
    batch_size: int = 4096,
    **kwargs: Any,
) -> list[pa.RecordBatch]:
    """Prefetch an entire catalog batch stream (tests / smoke / debug only)."""
    batches: list[pa.RecordBatch] = []
    async for batch in catalog.iter_batches(
        asset_ref=asset_ref,
        columns=columns,
        batch_size=batch_size,
        **kwargs,
    ):
        batches.append(batch)
    return batches

def to_torch_dataloader_from_catalog_batch_stream(
    stream: CatalogBatchStreamParams,
    *,
    sample_format: SampleFormat = "dict",
    feature_columns: list[str] | None = None,
    target_column: str | None = None,
    transform: Callable | None = None,
    row_preprocess: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    limit: int | None = None,
    batch_size: int = 32,
    num_workers: int = 0,
    execution_options: TorchExecutionOptions | None = None,
    queue_size: int = 4,
    stats: BatchStreamStats | None = None,
    **dataloader_kwargs: Any,
) -> Any:
    """Build a streaming ``DataLoader`` over a catalog batch stream (sync interop).

    Each full pass over the loader starts a new ``catalog.iter_batches`` stream
    via a background thread that spawns its own ``AsyncNexus``. Use
    :func:`catalog_batch_stream_for_epoch` or
    :meth:`CatalogBatchStreamParams.with_epoch_shuffle` for per-epoch streams.

    For async training with a caller-owned client, see
    :func:`async_iter_catalog_torch_samples` in
    ``archive/examples/nexus/catalog_torch_train.py`` (Layer 2 / power-user).

    ``row_preprocess`` + ``decode='row'`` is for string fixtures / custom transforms;
    numeric catalog columns should use ``TorchExecutionOptions(decode='vectorized')``.
    """
    opts = resolve_execution_options(execution_options)
    use_vectorized = opts.decode == "vectorized" and transform is None and batch_size >= 1
    dataset = _CatalogBatchStreamingDataset(
        stream,
        sample_format=sample_format,
        feature_columns=feature_columns,
        target_column=target_column,
        transform=transform,
        row_preprocess=row_preprocess,
        limit=limit,
        execution_options=execution_options,
        micro_batch_size=batch_size if use_vectorized else None,
        queue_size=queue_size,
        stats=stats,
    )
    return _build_catalog_batch_dataloader(
        dataset,
        batch_size=batch_size,
        use_vectorized=use_vectorized,
        num_workers=num_workers,
        **dataloader_kwargs,
    )

def to_torch_dataloader_from_prefetched_batches(
    record_batches: Sequence[pa.RecordBatch],
    *,
    sample_format: SampleFormat = "dict",
    feature_columns: list[str] | None = None,
    target_column: str | None = None,
    transform: Callable | None = None,
    row_preprocess: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    limit: int | None = None,
    batch_size: int = 32,
    num_workers: int = 0,
    execution_options: TorchExecutionOptions | None = None,
    stats: BatchStreamStats | None = None,
    **dataloader_kwargs: Any,
) -> Any:
    """Build a ``DataLoader`` from prefetched batches (tests / tiny fixtures only)."""
    opts = resolve_execution_options(execution_options)
    use_vectorized = opts.decode == "vectorized" and transform is None and batch_size >= 1
    dataset = _CatalogBatchPrefetchDataset(
        record_batches,
        sample_format=sample_format,
        feature_columns=feature_columns,
        target_column=target_column,
        transform=transform,
        row_preprocess=row_preprocess,
        limit=limit,
        execution_options=execution_options,
        micro_batch_size=batch_size if use_vectorized else None,
        stats=stats,
    )
    return _build_catalog_batch_dataloader(
        dataset,
        batch_size=batch_size,
        use_vectorized=use_vectorized,
        num_workers=num_workers,
        **dataloader_kwargs,
    )

def to_torch_iterable_dataset(
    collection: Any,
    *,
    sample_format: SampleFormat = "dict",
    feature_columns: list[str] | None = None,
    target_column: str | None = None,
    transform: Callable | None = None,
    limit: int | None = None,
    execution_options: TorchExecutionOptions | None = None,
    select: "ItemSelect | Mapping[str, Any] | None" = None,
    path: str | None = "/",
    shuffle: bool = False,
    client_shuffle_buffer: int | None = None,
    shuffle_seed: int | None = None,
) -> Any:
    """Build a torch ``IterableDataset`` over Parquet items in a Nexus collection.

    Server-side filtering: pass ``select`` and/or ``path`` to enumerate
    only matching items via ``AsyncCollection.iter_items``. Defaults preserve
    whole-collection traversal.

    When ``shuffle=True``, applies a bounded **client-side** shuffle buffer
    (default :data:`DEFAULT_CLIENT_SHUFFLE_BUFFER`). This is not catalog CRS
    epoch shuffle.
    """
    if isinstance(source := collection, DatasetSource):
        if source.readiness != "collection":
            raise TypeError(
                "collection torch adapter requires DatasetSource.from_collection(...); "
                "use miura.adapters.catalog.torch for catalog assets."
            )
        collection = source.collection
        select = source.select if select is None else select
        path = source.path if path == "/" else path

    dataset = _EpicoTorchIterableDataset(
        collection,
        sample_format=sample_format,
        feature_columns=feature_columns,
        target_column=target_column,
        transform=transform,
        limit=limit,
        execution_options=execution_options,
        batched_chunk_size=None,
        select=select,
        path=path,
    )
    if shuffle or client_shuffle_buffer is not None:
        buffer_size = client_shuffle_buffer if client_shuffle_buffer is not None else DEFAULT_CLIENT_SHUFFLE_BUFFER
        dataset = ClientShuffleIterableDataset(dataset, buffer_size=buffer_size, seed=shuffle_seed)
    return dataset

def to_torch_dataloader(
    source: Any,
    *,
    sample_format="dict",
    feature_columns=None,
    target_column=None,
    transform=None,
    row_preprocess=None,
    limit=None,
    batch_size=32,
    shuffle=False,
    num_workers=0,
    execution_options=None,
    queue_size=4,
    stats=None,
    select=None,
    path="/",
    client_shuffle_buffer=None,
    shuffle_seed=None,
    **dataloader_kwargs,
):
    """Build a torch ``DataLoader`` from a Nexus **collection** of Parquet items.

    Pass ``select=ItemSelect(split="train")`` for split-filtered training.
    ``shuffle=True`` uses a client-side buffer on the collection path.

    For catalog assets (Flight / CRS epoch shuffle), use
    :func:`miura.adapters.catalog.torch.to_torch_dataloader`.
    """
    if _is_catalog_source(source):
        warnings.warn(
            "Catalog assets should use miura.adapters.catalog.torch.to_torch_dataloader; "
            "routing catalog sources through miura.adapters.torch is deprecated.",
            DeprecationWarning,
            stacklevel=2,
        )
        from miura.adapters.catalog.torch import to_torch_dataloader as catalog_to_torch_dataloader

        return catalog_to_torch_dataloader(
            source,
            sample_format=sample_format,
            feature_columns=feature_columns,
            target_column=target_column,
            transform=transform,
            row_preprocess=row_preprocess,
            limit=limit,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            execution_options=execution_options,
            queue_size=queue_size,
            stats=stats,
            **dataloader_kwargs,
        )

    return _to_torch_dataloader_collection(
        source,
        sample_format=sample_format,
        feature_columns=feature_columns,
        target_column=target_column,
        transform=transform,
        limit=limit,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        execution_options=execution_options,
        select=select,
        path=path,
        client_shuffle_buffer=client_shuffle_buffer,
        shuffle_seed=shuffle_seed,
        **dataloader_kwargs,
    )

def to_dataloader(*args, **kwargs):
    """Compatibility alias; prefer :func:`to_torch_dataloader`."""
    return to_torch_dataloader(*args, **kwargs)