src/turn_detection/data/ingest.py

"""Lazy local-Parquet and Hugging Face dataset ingestion."""

from __future__ import annotations

import csv
import gzip
import json
from pathlib import Path
from typing import Any, Iterable, Iterator, Mapping, Sequence


class OptionalDependencyError(ImportError):
    """An optional data backend is required for the requested source."""


class DatasetReadError(RuntimeError):
    """A source exists but could not be decoded into records."""


def _open_text(path: Path):
    if path.suffix == ".gz":
        return gzip.open(path, "rt", encoding="utf-8", newline="")
    return path.open("r", encoding="utf-8", newline="")


def _with_provenance(row: Mapping[str, Any], source: str, index: int) -> dict[str, Any]:
    result = dict(row)
    result.setdefault("__source_file", source)
    result.setdefault("__source_row", index)
    return result


def _iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
    with _open_text(path) as handle:
        for index, line in enumerate(handle):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise DatasetReadError(f"invalid JSON at {path}:{index + 1}: {exc}") from exc
            if not isinstance(row, Mapping):
                raise DatasetReadError(f"expected an object at {path}:{index + 1}")
            yield _with_provenance(row, str(path), index)


def _iter_json(path: Path) -> Iterator[dict[str, Any]]:
    try:
        with _open_text(path) as handle:
            payload = json.load(handle)
    except (OSError, json.JSONDecodeError) as exc:
        raise DatasetReadError(f"cannot read JSON source {path}: {exc}") from exc
    if isinstance(payload, Mapping):
        for key in ("records", "data", "examples", "rows"):
            if key in payload:
                payload = payload[key]
                break
    if not isinstance(payload, list):
        raise DatasetReadError(f"JSON source {path} must contain a list of records")
    for index, row in enumerate(payload):
        if not isinstance(row, Mapping):
            raise DatasetReadError(f"expected an object at {path} record {index}")
        yield _with_provenance(row, str(path), index)


def _iter_csv(path: Path) -> Iterator[dict[str, Any]]:
    with _open_text(path) as handle:
        reader = csv.DictReader(handle)
        for index, row in enumerate(reader):
            yield _with_provenance(row, str(path), index)


def _import_pyarrow_parquet():
    try:
        import pyarrow.parquet as parquet  # type: ignore[import-not-found]
    except ImportError as exc:
        raise OptionalDependencyError(
            "reading local Parquet requires pyarrow; install the project's data dependencies"
        ) from exc
    return parquet


def iter_parquet_records(
    paths: Sequence[str | Path],
    *,
    batch_size: int = 256,
    columns: Sequence[str] | None = None,
) -> Iterator[dict[str, Any]]:
    """Stream rows from Parquet files without decoding Hugging Face audio."""

    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    parquet = _import_pyarrow_parquet()
    for raw_path in paths:
        path = Path(raw_path)
        try:
            parquet_file = parquet.ParquetFile(str(path))
            row_index = 0
            for batch in parquet_file.iter_batches(batch_size=batch_size, columns=columns):
                for row in batch.to_pylist():
                    if not isinstance(row, Mapping):
                        raise DatasetReadError(f"non-object Parquet row in {path}")
                    yield _with_provenance(row, str(path), row_index)
                    row_index += 1
        except OptionalDependencyError:
            raise
        except Exception as exc:
            raise DatasetReadError(f"cannot read Parquet source {path}: {exc}") from exc


def _local_files(source: Path) -> list[Path]:
    if source.is_file():
        return [source]
    if not source.exists():
        raise FileNotFoundError(f"dataset source does not exist: {source}")
    if not source.is_dir():
        raise DatasetReadError(f"dataset source is not a file or directory: {source}")
    supported = {".parquet", ".jsonl", ".ndjson", ".json", ".csv"}
    files = [path for path in source.rglob("*") if path.is_file() and path.suffix.lower() in supported]
    files.extend(
        path
        for path in source.rglob("*.jsonl.gz")
        if path.is_file() and path not in files
    )
    return sorted(files, key=lambda item: item.as_posix())


def iter_local_records(
    source: str | Path | Sequence[str | Path],
    *,
    batch_size: int = 256,
    columns: Sequence[str] | None = None,
) -> Iterator[dict[str, Any]]:
    """Iterate records from local Parquet, JSON(L), or CSV sources."""

    roots = [source] if isinstance(source, (str, Path)) else list(source)
    files: list[Path] = []
    for root in roots:
        files.extend(_local_files(Path(root)))
    files = sorted(set(files), key=lambda item: item.as_posix())
    if not files:
        raise DatasetReadError("no supported dataset files were found")

    parquet_paths = [path for path in files if path.suffix.lower() == ".parquet"]
    other_paths = [path for path in files if path.suffix.lower() != ".parquet"]
    if parquet_paths:
        yield from iter_parquet_records(parquet_paths, batch_size=batch_size, columns=columns)
    for path in other_paths:
        name = path.name.lower()
        if name.endswith((".jsonl", ".ndjson", ".jsonl.gz")):
            yield from _iter_jsonl(path)
        elif name.endswith(".json"):
            yield from _iter_json(path)
        elif name.endswith(".csv"):
            yield from _iter_csv(path)


def _normalize_hf_id(source: str) -> str:
    value = source
    if value.startswith("hf://datasets/"):
        value = value[len("hf://datasets/") :]
    elif value.startswith("hf://"):
        value = value[len("hf://") :]
    return value.strip("/")


def iter_hf_records(
    dataset_id: str,
    *,
    split: str = "train",
    revision: str | None = None,
    token: str | None = None,
    streaming: bool = True,
) -> Iterator[dict[str, Any]]:
    """Stream raw rows from a Hugging Face dataset with audio decoding disabled."""

    try:
        import datasets  # type: ignore[import-not-found]
    except ImportError as exc:
        raise OptionalDependencyError(
            "reading a Hugging Face dataset requires the 'datasets' package"
        ) from exc

    dataset_id = _normalize_hf_id(dataset_id)
    kwargs: dict[str, Any] = {"split": split, "streaming": streaming}
    if revision is not None:
        kwargs["revision"] = revision
    if token is not None:
        kwargs["token"] = token
    try:
        dataset = datasets.load_dataset(dataset_id, **kwargs)
    except TypeError as exc:
        # Compatibility with older datasets releases which used use_auth_token.
        if token is None or "token" not in str(exc):
            raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {exc}") from exc
        kwargs.pop("token", None)
        kwargs["use_auth_token"] = token
        try:
            dataset = datasets.load_dataset(dataset_id, **kwargs)
        except Exception as retry_exc:
            raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {retry_exc}") from retry_exc
    except Exception as exc:
        raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {exc}") from exc

    # ``decode=False`` keeps encoded bytes available and avoids torchcodec/ffmpeg.
    try:
        if "audio" in dataset.features:
            dataset = dataset.cast_column("audio", datasets.Audio(decode=False))
    except (AttributeError, TypeError, ValueError):
        # Some custom builders already return undecoded mappings or do not expose
        # cast_column in streaming mode.  Defer representation validation to audio.py.
        pass

    provenance = f"hf://datasets/{dataset_id}@{revision or 'default'}:{split}"
    for index, row in enumerate(dataset):
        if not isinstance(row, Mapping):
            raise DatasetReadError(f"Hugging Face dataset {dataset_id} yielded a non-object row")
        yield _with_provenance(row, provenance, index)


def _looks_like_hf_id(source: str) -> bool:
    if source.startswith("hf://"):
        return True
    return source.count("/") == 1 and not source.startswith(("./", "../", "/"))


def iter_records(
    source: str | Path | Sequence[str | Path],
    *,
    split: str = "train",
    revision: str | None = None,
    token: str | None = None,
    streaming: bool = True,
    batch_size: int = 256,
    columns: Sequence[str] | None = None,
    limit: int | None = None,
) -> Iterator[dict[str, Any]]:
    """Iterate local or Hugging Face rows through one stable entry point."""

    if limit is not None and limit < 0:
        raise ValueError("limit cannot be negative")
    if isinstance(source, str) and _looks_like_hf_id(source) and not Path(source).exists():
        iterator: Iterable[dict[str, Any]] = iter_hf_records(
            source,
            split=split,
            revision=revision,
            token=token,
            streaming=streaming,
        )
    else:
        iterator = iter_local_records(source, batch_size=batch_size, columns=columns)
    for index, row in enumerate(iterator):
        if limit is not None and index >= limit:
            break
        yield row