"""Fetch a dataset from its original source and harmonize it locally.
PaMIR ships code and a recipe, not data. ``download(id)`` fetches the raw
file from the dataset's original source into a local cache, harmonizes it
(see :mod:`pamir.harmonize`), writes a parquet next to it, and validates the
result against the shipped ``expected.json`` schema. ``load_dataset`` calls
this automatically on first use.
Sources and their requirements:
- ``kaggle`` — needs ``kagglehub`` and Kaggle credentials
(``KAGGLE_USERNAME`` + ``KAGGLE_KEY``, or ``~/.kaggle/kaggle.json``); some
competition datasets require accepting their rules on the Kaggle website once.
- ``uci_zip`` — public download, no credentials.
- ``github_raw`` — public download, no credentials.
- ``hf`` — needs ``huggingface_hub``.
Install the optional fetch dependencies with ``pip install pamir-credit[data]``.
"""
import io
import json
import os
import zipfile
from pathlib import Path
from typing import Dict, Optional
from urllib.request import Request, urlopen
import pandas as pd
from pamir.catalog import dataset_info, open_datasets
from pamir.harmonize import harmonize
_EXPECTED = json.loads(
(Path(__file__).parent / "data" / "expected.json").read_text(encoding="utf-8")
)
[docs]
def cache_dir() -> Path:
"""Directory where fetched raw files and harmonized parquet are cached."""
env = os.environ.get("PAMIR_CACHE")
if env:
base = Path(env)
else:
try:
from platformdirs import user_cache_dir
base = Path(user_cache_dir("pamir"))
except Exception:
base = Path.home() / ".cache" / "pamir"
base.mkdir(parents=True, exist_ok=True)
return base
def _http_get(url: str) -> bytes:
req = Request(url, headers={"User-Agent": "Mozilla/5.0 (compatible; pamir-credit)"})
with urlopen(req, timeout=120) as r: # noqa: S310 (trusted, https)
return r.read()
def _extract_from_zip(blob: bytes, fname: str, _depth: int = 0):
"""Return the bytes of ``fname`` from a zip, recursing one level into
nested zips (the PAKDD archive nests the modeling data inside another zip)."""
with zipfile.ZipFile(io.BytesIO(blob)) as z:
member = next((m for m in z.namelist() if m.endswith(fname)), None)
if member is not None:
return z.read(member)
if _depth == 0:
for m in z.namelist():
if m.lower().endswith(".zip"):
got = _extract_from_zip(z.read(m), fname, _depth + 1)
if got is not None:
return got
return None
def _fetch_raw(spec: Dict, raw_dir: Path) -> Path:
"""Fetch the raw source file into ``raw_dir`` and return its path."""
dl = spec["download"]
kind, loc, fname = dl["kind"], dl["locator"], dl["file"]
raw_dir.mkdir(parents=True, exist_ok=True)
target = raw_dir / fname
if kind == "github_raw":
target.write_bytes(_http_get(loc))
return target
if kind in ("uci_zip", "github_zip"):
data = _extract_from_zip(_http_get(loc), fname)
if data is None:
raise FileNotFoundError(f"{fname} not found in zip for '{spec['id']}'")
target.write_bytes(data)
return target
if kind == "kaggle_competition":
try:
import kagglehub
except ImportError as e:
raise ImportError(
"kaggle competition source needs `kagglehub` — pip install pamir-credit[data]"
) from e
# Single-file download (path=) — avoids pulling the whole ~700MB
# competition, and kagglehub gives a clear message if the rules were not
# accepted (a bare 401 otherwise). Rules must be accepted once at
# https://www.kaggle.com/competitions/<comp>/rules
try:
got = Path(kagglehub.competition_download(loc, path=fname))
except TypeError: # older kagglehub without path=
got = Path(kagglehub.competition_download(loc))
if got.is_file():
return got
hit = next((p for p in got.rglob("*") if p.name == fname), None)
if hit is None:
raise FileNotFoundError(
f"{fname} not found in competition '{loc}' (under {got})")
return hit
if kind == "kaggle":
try:
import kagglehub
except ImportError as e:
raise ImportError(
"kaggle source needs `kagglehub` — pip install pamir-credit[data]"
) from e
path = Path(kagglehub.dataset_download(loc))
hit = next((p for p in path.rglob("*") if p.name == fname), None)
if hit is None:
raise FileNotFoundError(
f"{fname} not found in Kaggle dataset '{loc}' for '{spec['id']}'")
return hit
if kind == "hf":
try:
from huggingface_hub import hf_hub_download
except ImportError as e:
raise ImportError(
"hf source needs `huggingface_hub` — pip install pamir-credit[data]"
) from e
return Path(hf_hub_download(repo_id=loc, filename=fname, repo_type="dataset"))
raise ValueError(f"unknown download.kind {kind!r} for '{spec['id']}'")
def _validate(dataset_id: str, df: pd.DataFrame) -> list:
"""Compare a harmonized table against expected.json; return warnings."""
exp = _EXPECTED.get(dataset_id)
warn = []
if exp is None:
return ["no expected.json entry — result not validated"]
y = df["__target__"]
n, dr = len(df), float(y.mean())
nfeat = df.shape[1] - 1
if y.sum() == 0:
warn.append("target parsed as all-zero (harmonization/source problem)")
if nfeat != exp["n_features"]:
got = set(df.columns) - {"__target__"}
want = set(exp["columns"])
warn.append(f"features {nfeat} != expected {exp['n_features']} "
f"(missing {sorted(want - got)[:6]}, extra {sorted(got - want)[:6]})")
if abs(n - exp["n_rows"]) / max(exp["n_rows"], 1) > 0.02:
warn.append(f"rows {n:,} != expected {exp['n_rows']:,} "
"(source snapshot may have changed since the benchmark was built)")
if abs(dr - exp["DR"]) > 0.02:
warn.append(f"DR {dr:.3f} != expected {exp['DR']:.3f}")
return warn
[docs]
def download(dataset_id: str, force: bool = False, quiet: bool = False) -> Path:
"""Fetch and harmonize one dataset; return the path to the cached parquet.
Parameters
----------
dataset_id : str
A PaMIR dataset id (see :func:`pamir.list_datasets`).
force : bool
Re-fetch and re-harmonize even if a cached parquet exists.
quiet : bool
Suppress progress / validation messages.
Returns
-------
Path to ``<cache>/<dataset_id>.parquet``.
"""
spec = dataset_info(dataset_id)
out = cache_dir() / f"{dataset_id}.parquet"
if out.exists() and not force:
return out
def say(msg):
if not quiet:
print(f"[pamir] {dataset_id}: {msg}", flush=True)
say(f"fetching from {spec['download']['kind']} …")
raw = _fetch_raw(spec, cache_dir() / "raw" / dataset_id)
say("harmonizing …")
df = harmonize(raw, spec)
for w in _validate(dataset_id, df):
say(f"WARNING: {w}")
df.to_parquet(out, index=False)
say(f"cached {len(df):,} rows, {df.shape[1] - 1} features -> {out}")
return out
def download_open(force: bool = False, quiet: bool = False) -> list:
"""Download every dataset whose source needs no credentials.
A zero-setup starter set (UCI / GitHub / HuggingFace) — no
Kaggle account required. Returns the list of cached parquet paths.
"""
paths = []
for ds in open_datasets():
paths.append(download(ds, force=force, quiet=quiet))
return paths