| """Rows read straight from a dataset's files, for when the dataset viewer can't (its conversion failed, or it's off): |
| JSON Lines, JSON (a list, or a dict of lists), parquet and CSV. |
| |
| Parquet is read a row group at a time over HTTP ranges, so a big file costs only the groups a page needs. A JSON |
| Lines file is downloaded once (to the Hub cache) and indexed by line offsets, so any row is one seek. Search and |
| filters scan the first rows only (up to SCAN), and say so. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import posixpath |
| import re |
| import threading |
| import time |
| from array import array |
| from typing import Any |
|
|
| from .. import catalog, config |
|
|
| EXT = re.compile(r"\.(jsonl|ndjson|json|parquet|csv|tsv)(\.gz)?$", re.I) |
| SPLIT = re.compile(r"(^|[/_.\-])(train|training|test|testing|eval|evaluation|valid|validation|dev|val)([/_.\-]|$)", re.I) |
| MAX_FILE = 600 * 2**20 |
| SCAN = 20_000 |
| _lock = threading.Lock() |
| _lines: dict[str, tuple[str, array]] = {} |
|
|
|
|
| class DirectError(Exception): |
| def __init__(self, message: str, status: int = 502): |
| super().__init__(message) |
| self.status = status |
|
|
|
|
| HEAD = 512 * 2**10 |
| _fetching: dict[str, dict[str, Any]] = {} |
|
|
|
|
| def _head(spec: str, sha: str, path: str, token: str | None) -> tuple[list[dict[str, Any]], bool]: |
| """The rows in the first HEAD bytes of a JSON Lines file (one range request), and whether that was all of it.""" |
| import httpx |
| from urllib.parse import quote |
|
|
| def fetch(): |
| r = httpx.get(f"https://huggingface.co/datasets/{spec}/resolve/{sha}/{quote(path)}", follow_redirects=True, timeout=httpx.Timeout(45, connect=10), |
| headers={"Range": f"bytes=0-{HEAD - 1}", **({"Authorization": f"Bearer {token}"} if token else {})}) |
| if r.status_code not in (200, 206): |
| raise DirectError(f"couldn't read {path} (HTTP {r.status_code})") |
| data = r.content |
| whole = r.status_code == 200 or len(data) < HEAD |
| lines = data.split(b"\n") |
| if not whole: |
| lines = lines[:-1] |
| out = [] |
| for ln in lines: |
| if ln.strip(): |
| try: |
| v = json.loads(ln) |
| except ValueError: |
| continue |
| out.append(v if isinstance(v, dict) else {"value": v}) |
| return out, whole |
|
|
| return catalog._cached(("direct-head", spec, sha, path), 3600, fetch) |
|
|
|
|
| def _background(spec: str, sha: str, path: str, size: int | None, token: str | None) -> dict[str, Any]: |
| """The whole file, downloaded in the background once; its state now.""" |
| key = f"{spec}@{sha}/{path}" |
| with _lock: |
| st = _fetching.get(key) |
| if st and (st["state"] != "error" or time.time() - st.get("at", 0) < 120): |
| return st |
| st = _fetching[key] = {"state": "running", "at": time.time()} |
|
|
| def run(): |
| try: |
| st["local"] = _local(spec, sha, path, size, token) |
| _line_index(st["local"]) |
| st["state"] = "done" |
| except Exception as e: |
| st.update(state="error", error=str(e)[:200], at=time.time()) |
|
|
| threading.Thread(target=run, daemon=True, name=f"direct-{path}").start() |
| return st |
|
|
|
|
| def data_files(spec: str, meta: dict[str, Any], token: str | None) -> dict[str, dict[str, list[str]]]: |
| """{config: {split: [paths]}} from the card's `configs`, or the data files the repository holds.""" |
| from fnmatch import fnmatch |
|
|
| def listing(): |
| out = [] |
| for n, e in enumerate(catalog._api(token).list_repo_tree(spec, repo_type="dataset", revision=meta["sha"], recursive=True)): |
| if hasattr(e, "size") and EXT.search(e.path) and not posixpath.basename(e.path).startswith("."): |
| out.append((e.path, e.size)) |
| if n > 50_000: |
| break |
| return out |
|
|
| files = catalog._cached(("data-files", spec, meta["sha"]), 3600, listing) |
| names = [p for p, _ in files] |
| configs: dict[str, dict[str, list[str]]] = {} |
|
|
| def card_data(): |
| info = catalog._api(token).dataset_info(spec, revision=meta["sha"], expand=["cardData"]) |
| return info.card_data.to_dict() if info.card_data else {} |
|
|
| card = catalog._cached(("card", spec, meta["sha"]), 3600, card_data) |
| for c in (card.get("configs") or []) if isinstance(card, dict) else []: |
| df = c.get("data_files") |
| name = c.get("config_name") or "default" |
| specs = [{"split": "train", "path": df}] if isinstance(df, (str, list)) and not (isinstance(df, list) and df and isinstance(df[0], dict)) else (df or []) |
| for d in specs: |
| pats = d["path"] if isinstance(d.get("path"), list) else [d.get("path")] |
| hits = sorted(p for p in names if any(pat and fnmatch(p, pat) for pat in pats)) |
| if hits: |
| configs.setdefault(name, {})[d.get("split") or "train"] = hits |
| if not configs and names: |
| by: dict[str, list[str]] = {} |
| for p in sorted(names): |
| m = SPLIT.search(p) |
| split = {"training": "train", "testing": "test", "evaluation": "eval", "valid": "validation", "val": "validation"}.get( |
| m.group(2).lower(), m.group(2).lower()) if m else "train" |
| by.setdefault(split, []).append(p) |
| configs["default"] = by |
| return configs |
|
|
|
|
| def _local(spec: str, sha: str, path: str, size: int | None, token: str | None) -> str: |
| from huggingface_hub import hf_hub_download |
|
|
| if size is not None and size > MAX_FILE: |
| raise DirectError(f"{path} is {size / 2**20:.0f} MB: too large to read here") |
| return hf_hub_download(spec, path, repo_type="dataset", revision=sha, token=token or False, |
| cache_dir=str(config.CACHE_DIR / "direct")) |
|
|
|
|
| def _line_index(local: str) -> array: |
| with _lock: |
| hit = _lines.get(local) |
| if hit: |
| return hit[1] |
| offs = array("Q") |
| with open(local, "rb") as f: |
| pos = 0 |
| for line in f: |
| if line.strip(): |
| offs.append(pos) |
| pos += len(line) |
| with _lock: |
| _lines[local] = (local, offs) |
| while len(_lines) > 12: |
| _lines.pop(next(iter(_lines))) |
| return offs |
|
|
|
|
| class Source: |
| """One split's rows, across its files, in order.""" |
|
|
| def __init__(self, spec: str, meta: dict[str, Any], paths: list[str], token: str | None): |
| self.spec, self.meta, self.paths, self.token = spec, meta, paths[:200], token |
| self._counts: list[int] | None = None |
|
|
| def _size(self, path: str) -> int | None: |
| files = catalog._cached(("data-files", self.spec, self.meta["sha"]), 3600, lambda: []) |
| return next((s for p, s in files if p == path), None) |
|
|
| def _parquet(self, path: str): |
| import pyarrow.parquet as pq |
| from huggingface_hub import HfFileSystem |
|
|
| fs = HfFileSystem(token=self.token or False) |
| return pq.ParquetFile(fs.open(f"datasets/{self.spec}@{self.meta['sha']}/{path}", "rb", block_size=2**20)) |
|
|
| def count(self, path: str) -> int: |
| def fetch(): |
| low = path.lower() |
| if low.endswith(".parquet"): |
| return self._parquet(path).metadata.num_rows |
| if low.endswith((".jsonl", ".ndjson")): |
| st = _background(self.spec, self.meta["sha"], path, self._size(path), self.token) |
| if st["state"] != "done": |
| rows, whole = _head(self.spec, self.meta["sha"], path, self.token) |
| return len(rows) if whole else None |
| return len(_line_index(st["local"])) |
| return len(self._whole(path)) |
|
|
| n = catalog._cached(("direct-count", self.spec, self.meta["sha"], path), 3600, fetch) |
| if n is None: |
| catalog._memo.pop(("direct-count", self.spec, self.meta["sha"], path), None) |
| return n |
|
|
| def _whole(self, path: str) -> list[dict[str, Any]]: |
| def fetch(): |
| local = _local(self.spec, self.meta["sha"], path, self._size(path), self.token) |
| low = path.lower() |
| if low.endswith((".csv", ".tsv")): |
| import csv |
|
|
| with open(local, newline="", encoding="utf-8", errors="replace") as f: |
| return list(csv.DictReader(f, delimiter="\t" if low.endswith(".tsv") else ",")) |
| with open(local, encoding="utf-8", errors="replace") as f: |
| doc = json.load(f) |
| if isinstance(doc, dict): |
| lists = [v for v in doc.values() if isinstance(v, list)] |
| if len(doc) == 1 and lists: |
| doc = lists[0] |
| elif lists and all(isinstance(v, list) and len(v) == len(lists[0]) for v in doc.values()): |
| doc = [dict(zip(doc, vals)) for vals in zip(*doc.values())] |
| else: |
| doc = [doc] |
| return [r if isinstance(r, dict) else {"value": r} for r in doc] if isinstance(doc, list) else [] |
|
|
| return catalog._cached(("direct-whole", self.spec, self.meta["sha"], path), 1800, fetch, big=True) |
|
|
| def total(self) -> int | None: |
| """Rows in the split; None while a file is still being read (the first rows are served meanwhile).""" |
| if self._counts is None: |
| self._counts = [self.count(p) for p in self.paths] |
| return None if any(n is None for n in self._counts) else sum(self._counts) |
|
|
| def read(self, offset: int, length: int) -> list[tuple[int, dict[str, Any]]]: |
| self.total() |
| if self._counts and self._counts[0] is None: |
| rows, _ = _head(self.spec, self.meta["sha"], self.paths[0], self.token) |
| if offset + length <= len(rows) or offset < len(rows): |
| return [(offset + k, r) for k, r in enumerate(rows[offset:offset + length])] |
| raise DirectError("this file is still being read; these rows show in a minute", 409) |
| out: list[tuple[int, dict[str, Any]]] = [] |
| base = 0 |
| for path, n in zip(self.paths, self._counts or []): |
| if n is None: |
| raise DirectError("this file is still being read; these rows show in a minute", 409) |
| if offset + length <= base: |
| break |
| if offset < base + n: |
| lo, hi = max(0, offset - base), min(n, offset + length - base) |
| out += [(base + lo + k, r) for k, r in enumerate(self._slice(path, lo, hi))] |
| base += n |
| return out |
|
|
| def _slice(self, path: str, lo: int, hi: int) -> list[dict[str, Any]]: |
| low = path.lower() |
| if low.endswith(".parquet"): |
| pf = self._parquet(path) |
| rows, start = [], 0 |
| for g in range(pf.metadata.num_row_groups): |
| n = pf.metadata.row_group(g).num_rows |
| if start + n > lo and start < hi: |
| tbl = pf.read_row_group(g).to_pylist() |
| rows += tbl[max(0, lo - start):hi - start] |
| start += n |
| if start >= hi: |
| break |
| return [_plain(r) for r in rows] |
| if low.endswith((".jsonl", ".ndjson")): |
| st = _background(self.spec, self.meta["sha"], path, self._size(path), self.token) |
| if st["state"] != "done": |
| raise DirectError("this file is still being read; these rows show in a minute", 409) |
| local = st["local"] |
| offs = _line_index(local) |
| out = [] |
| with open(local, "rb") as f: |
| for k in range(lo, min(hi, len(offs))): |
| f.seek(offs[k]) |
| try: |
| r = json.loads(f.readline()) |
| except ValueError: |
| r = {"_unreadable_line": k} |
| out.append(r if isinstance(r, dict) else {"value": r}) |
| return out |
| return self._whole(path)[lo:hi] |
|
|
|
|
| def _plain(r: dict[str, Any]) -> dict[str, Any]: |
| """Parquet values the page can take: bytes as base64 (a packed task), the rest as they are.""" |
| import base64 |
|
|
| out = {} |
| for k, v in r.items(): |
| out[k] = base64.b64encode(v).decode() if isinstance(v, (bytes, bytearray)) else v |
| return out |
|
|
|
|
| def scan(src: Source, match, limit: int = SCAN) -> tuple[list[tuple[int, dict[str, Any]]], bool]: |
| """Rows that match, from the first `limit` rows; and whether it stopped early.""" |
| hits = [] |
| total = src.total() |
| if total is None: |
| raise DirectError("search works once the whole file is read; try again in a minute", 409) |
| for start in range(0, min(total, limit), 500): |
| for i, r in src.read(start, 500): |
| if match(r): |
| hits.append((i, r)) |
| return hits, total > limit |
|
|
