| """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 |
| 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 |
| 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: |
| |
| 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 |
|
|
| |
| try: |
| if "audio" in dataset.features: |
| dataset = dataset.cast_column("audio", datasets.Audio(decode=False)) |
| except (AttributeError, TypeError, ValueError): |
| |
| |
| 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 |
|
|
