"""Lazy resolution of manifest provenance back to raw records. Parquet has row-group rather than true single-row random access. The resolver therefore keeps a small LRU of row groups. Iterating a manifest in its original audit order is efficient; arbitrary random access is intended for evaluation and inspection rather than full-data training shuffles. """ from __future__ import annotations from bisect import bisect_right from collections import OrderedDict from pathlib import Path from typing import Any, Iterable, Iterator, Mapping, Sequence from .ingest import DatasetReadError, _import_pyarrow_parquet, iter_local_records class ManifestResolutionError(DatasetReadError): """A manifest row cannot be resolved to its source record.""" class ManifestRecordResolver: """Resolve ``source_file``/``source_row`` provenance with bounded caching.""" def __init__( self, *, source_root: str | Path | None = None, max_cached_row_groups: int = 1, ) -> None: if max_cached_row_groups < 0: raise ValueError("max_cached_row_groups cannot be negative") self.source_root = Path(source_root) if source_root is not None else None self.max_cached_row_groups = max_cached_row_groups self._parquet_files: dict[Path, Any] = {} self._row_group_starts: dict[Path, list[int]] = {} self._row_group_cache: OrderedDict[tuple[Path, int, tuple[str, ...] | None], Any] = OrderedDict() self._basename_cache: dict[str, Path] = {} def clear(self) -> None: """Drop cached Parquet handles and row groups.""" self._parquet_files.clear() self._row_group_starts.clear() self._row_group_cache.clear() def _resolve_source_path(self, source_file: Any) -> Path: if not source_file: raise ManifestResolutionError("manifest row has no source_file") source_text = str(source_file) if source_text.startswith("hf://"): raise ManifestResolutionError( "remote Hugging Face provenance is not random-accessible; audit a downloaded snapshot instead" ) original = Path(source_text).expanduser() candidates = [original] if self.source_root is not None: candidates.extend((self.source_root / original, self.source_root / original.name)) for candidate in candidates: if candidate.is_file(): return candidate.resolve() if self.source_root is not None: cached = self._basename_cache.get(original.name) if cached is not None: return cached matches = sorted(self.source_root.rglob(original.name)) if len(matches) == 1: resolved = matches[0].resolve() self._basename_cache[original.name] = resolved return resolved if len(matches) > 1: raise ManifestResolutionError( f"source basename {original.name!r} is ambiguous below {self.source_root}" ) raise ManifestResolutionError(f"source file does not exist: {source_text}") def _parquet_handle(self, path: Path): handle = self._parquet_files.get(path) if handle is None: parquet = _import_pyarrow_parquet() try: handle = parquet.ParquetFile(str(path)) except Exception as exc: raise ManifestResolutionError(f"cannot open Parquet source {path}: {exc}") from exc self._parquet_files[path] = handle starts = [0] running = 0 for index in range(handle.metadata.num_row_groups): running += handle.metadata.row_group(index).num_rows starts.append(running) self._row_group_starts[path] = starts return handle def _read_parquet_row( self, path: Path, row_index: int, columns: Sequence[str] | None, ) -> dict[str, Any]: handle = self._parquet_handle(path) starts = self._row_group_starts[path] if row_index < 0 or row_index >= starts[-1]: raise ManifestResolutionError( f"source_row {row_index} is outside [0, {starts[-1]}) for {path}" ) row_group = bisect_right(starts, row_index) - 1 column_key = tuple(columns) if columns is not None else None cache_key = (path, row_group, column_key) table = self._row_group_cache.get(cache_key) if table is None: try: table = handle.read_row_group(row_group, columns=columns) except Exception as exc: raise ManifestResolutionError( f"cannot read row group {row_group} from {path}: {exc}" ) from exc if self.max_cached_row_groups: self._row_group_cache[cache_key] = table self._row_group_cache.move_to_end(cache_key) while len(self._row_group_cache) > self.max_cached_row_groups: self._row_group_cache.popitem(last=False) else: self._row_group_cache.move_to_end(cache_key) offset = row_index - starts[row_group] rows = table.slice(offset, 1).to_pylist() if not rows: raise ManifestResolutionError(f"failed to resolve row {row_index} from {path}") return dict(rows[0]) def resolve( self, manifest_row: Mapping[str, Any], *, columns: Sequence[str] | None = None, ) -> dict[str, Any]: """Load one raw source record referenced by a manifest row.""" path = self._resolve_source_path(manifest_row.get("source_file")) try: row_index = int(manifest_row.get("source_row")) except (TypeError, ValueError) as exc: raise ManifestResolutionError("manifest source_row is missing or invalid") from exc if path.suffix.lower() == ".parquet": record = self._read_parquet_row(path, row_index, columns) else: record = {} for candidate in iter_local_records(path, columns=columns): if int(candidate.get("__source_row", -1)) == row_index: record = candidate break if not record: raise ManifestResolutionError(f"failed to resolve row {row_index} from {path}") record["__source_file"] = str(path) record["__source_row"] = row_index return record def resolve_audio(self, manifest_row: Mapping[str, Any]) -> Any: """Resolve only the raw ``audio`` value for a manifest row.""" record = self.resolve(manifest_row, columns=("audio",)) if "audio" not in record: raise ManifestResolutionError("source record has no audio column") return record["audio"] def iter_manifest_records( rows: Iterable[Mapping[str, Any]], *, source_root: str | Path | None = None, columns: Sequence[str] | None = None, max_cached_row_groups: int = 1, ) -> Iterator[dict[str, Any]]: """Resolve manifest rows lazily; source-order manifests reuse row groups.""" resolver = ManifestRecordResolver( source_root=source_root, max_cached_row_groups=max_cached_row_groups, ) try: for row in rows: yield resolver.resolve(row, columns=columns) finally: resolver.clear()