Spaces:
Sleeping
Sleeping
| """Anonymous public-S3 Parquet scans with pinned manifests and resumable output. | |
| Run this batch code on the data/compute host. DuckDB reads selected columns over | |
| HTTPS range requests; it never downloads the full snapshot onto a browser client. | |
| """ | |
| import hashlib | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| from collections.abc import Callable | |
| from pathlib import Path | |
| from urllib.parse import urlparse | |
| import httpx | |
| from app.ingest.fetch import atomic_json, request_json | |
| from app.ingest.scope import PROFILES, SCOPE_VERSION, coarse_predicate, matches_work | |
| PUBLIC_MANIFEST = "https://openalex.s3.amazonaws.com/data/parquet/works/manifest.json" | |
| PUBLIC_HOSTS = {"openalex.s3.amazonaws.com", "openalex.s3.us-east-1.amazonaws.com"} | |
| PROJECTED_FIELDS = ( | |
| "id", | |
| "title", | |
| "display_name", | |
| "abstract_inverted_index", | |
| "topics", | |
| "type", | |
| "publication_year", | |
| "publication_date", | |
| "ids", | |
| "pmid", | |
| "indexed_in", | |
| "locations", | |
| "doi", | |
| "referenced_works", | |
| "is_retracted", | |
| "cited_by_count", | |
| "primary_location", | |
| "authorships", | |
| "updated_date", | |
| ) | |
| def fingerprint(value) -> str: | |
| return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() | |
| def public_url(value: str) -> str: | |
| """Force anonymous HTTPS to the public bucket, never a signed/paid/API endpoint.""" | |
| parsed = urlparse(value) | |
| if parsed.scheme == "s3" and parsed.netloc == "openalex": | |
| value = "https://openalex.s3.amazonaws.com" + parsed.path | |
| parsed = urlparse(value) | |
| if ( | |
| parsed.scheme != "https" | |
| or parsed.netloc not in PUBLIC_HOSTS | |
| or not parsed.path.startswith("/data/parquet/works/") | |
| or parsed.query | |
| or parsed.fragment | |
| or "/../" in parsed.path | |
| ): | |
| raise ValueError("Expected an unsigned public OpenAlex works Parquet S3 URL") | |
| return value | |
| async def load_manifest(location: str, *, client: httpx.AsyncClient | None = None) -> dict: | |
| if "://" in location: | |
| url = public_url(location) | |
| owned = client is None | |
| client = client or httpx.AsyncClient(timeout=60) | |
| try: | |
| data = await request_json(client, url, {}) | |
| finally: | |
| if owned: | |
| await client.aclose() | |
| else: | |
| data = json.loads(Path(location).read_text()) | |
| if not isinstance(data, dict) or not isinstance(data.get("files"), list): | |
| raise ValueError("Expected an OpenAlex works manifest containing a files array") | |
| if data.get("entity", "works") != "works" or data.get("format", "parquet") != "parquet": | |
| raise ValueError("Only the works Parquet manifest is supported") | |
| return data | |
| def plan_manifest( | |
| manifest: dict, | |
| *, | |
| max_files: int | None = None, | |
| largest_first: bool = False, | |
| skip_files: int = 0, | |
| ) -> dict: | |
| """Select manifest parts to scan. | |
| Parts are listed in update-date order and vary in size by four orders of | |
| magnitude, while per-part scan time is dominated by fixed overhead. For an | |
| explicitly partial run, ``largest_first`` therefore covers far more of the | |
| corpus per unit of time. It never changes which parts exist, only their | |
| order, so a full run is unaffected. | |
| ``skip_files`` continues after an earlier ``max_files`` run in a new data | |
| directory. The manifest fingerprint stays that of the whole release, and a | |
| run that skips parts never counts as having selected all of them. | |
| """ | |
| if max_files is not None and max_files <= 0: | |
| raise ValueError("max-files must be positive") | |
| if skip_files < 0: | |
| raise ValueError("skip-files must not be negative") | |
| parts = [] | |
| seen = set() | |
| for entry in manifest["files"]: | |
| location = entry.get("url", "") | |
| if "://" in location: | |
| location = public_url(location) | |
| elif location: | |
| location = str(Path(location).resolve()) | |
| size = int(entry.get("size_bytes") or entry.get("meta", {}).get("content_length") or 0) | |
| if not location.endswith(".parquet") or size <= 0 or location in seen: | |
| raise ValueError("Manifest requires unique Parquet paths and positive file sizes") | |
| seen.add(location) | |
| parts.append({"url": location, "size_bytes": size}) | |
| if not parts: | |
| raise ValueError("Manifest has no Parquet files") | |
| ordered = sorted(parts, key=lambda part: -part["size_bytes"]) if largest_first else parts | |
| remaining = ordered[skip_files:] | |
| if not remaining: | |
| raise ValueError("skip-files leaves no manifest parts to scan") | |
| selected = remaining[:max_files] if max_files is not None else remaining | |
| total_bytes = sum(part["size_bytes"] for part in parts) | |
| budgeted = sum(part["size_bytes"] for part in selected) | |
| return { | |
| "manifest_sha256": fingerprint(manifest), | |
| "snapshot_date": manifest.get("date"), | |
| "total_parts": len(parts), | |
| "selected_parts": len(selected), | |
| "skipped_parts": skip_files, | |
| "selection_order": "largest_first" if largest_first else "manifest", | |
| "physical_bytes_budgeted": budgeted, | |
| "selected_bytes_fraction": round(budgeted / total_bytes, 6) if total_bytes else 0.0, | |
| "all_parts_selected": len(parts) == len(selected), | |
| "files": selected, | |
| "note": "Physical file sizes are a conservative scan budget, not bytes actually transferred. Partitions are update dates, not topics. A partial selection is a subset of update partitions, not a random sample of the corpus.", | |
| } | |
| def _connection(extension_directory: Path): | |
| import duckdb | |
| extension_directory.mkdir(parents=True, exist_ok=True) | |
| connection = duckdb.connect(config={"extension_directory": str(extension_directory.resolve())}) | |
| connection.execute("SET memory_limit='1GB'") | |
| connection.execute("SET TimeZone='UTC'") | |
| # Stable row order lets a batch checkpoint resume within a pinned part. | |
| connection.execute("SET threads=1") | |
| connection.execute("SET preserve_insertion_order=true") | |
| return connection | |
| def scan_snapshot( | |
| files: list[str], | |
| output: Path, | |
| *, | |
| max_bytes: int = 5_000_000_000, | |
| limit: int | None = None, | |
| topic: str = "", | |
| profile: str = "all", | |
| remote_sizes: dict[str, int] | None = None, | |
| resume: bool = False, | |
| snapshot_id: str = "", | |
| snapshot_date: str | None = None, | |
| all_parts_selected: bool = False, | |
| work_ids: set[str] | None = None, | |
| batch_size: int = 1000, | |
| min_free_bytes: int = 1_000_000_000, | |
| progress: Callable[[dict], None] | None = None, | |
| ) -> dict: | |
| if ( | |
| not files | |
| or max_bytes <= 0 | |
| or batch_size <= 0 | |
| or min_free_bytes < 0 | |
| or (limit is not None and limit <= 0) | |
| ): | |
| raise ValueError("Provide files, positive budgets and a positive optional limit") | |
| if profile not in PROFILES: | |
| raise ValueError("Unknown corpus profile") | |
| if work_ids is not None and ( | |
| not work_ids or any(not re.fullmatch(r"W\d+", x) for x in work_ids) | |
| ): | |
| raise ValueError("Reference IDs must be a nonempty set of W identifiers") | |
| sources, identities, total = [], [], 0 | |
| for filename in files: | |
| if any(character in filename for character in "*?["): | |
| raise ValueError("Snapshot scan requires explicit files, not globs or signed URLs") | |
| if "://" in filename: | |
| source = public_url(filename) | |
| size = (remote_sizes or {}).get(filename) or (remote_sizes or {}).get(source) | |
| if not size or size <= 0: | |
| raise ValueError("Remote snapshot files require manifest size_bytes before reading") | |
| identity = [source, size] | |
| else: | |
| path = Path(filename).resolve() | |
| source, size = str(path), path.stat().st_size | |
| identity = [source, size, path.stat().st_mtime_ns] | |
| sources.append(source) | |
| identities.append(identity) | |
| total += size | |
| if total > max_bytes: | |
| raise ValueError( | |
| f"Snapshot files total {total} bytes, exceeding budget {max_bytes}; inspect the plan on the batch host" | |
| ) | |
| signature = fingerprint( | |
| [identities, snapshot_id, profile, topic, SCOPE_VERSION, sorted(work_ids or [])] | |
| ) | |
| checkpoint = output.with_suffix(output.suffix + ".checkpoint.json") | |
| state = { | |
| "signature": signature, | |
| "snapshot_id": snapshot_id, | |
| "snapshot_date": snapshot_date, | |
| "scope": "reference_backfill" if work_ids is not None else profile, | |
| "scope_version": SCOPE_VERSION, | |
| "records": 0, | |
| "metadata_only": 0, | |
| "bytes": 0, | |
| "part": 0, | |
| "part_rows": 0, | |
| "selected_parts": len(sources), | |
| "physical_bytes_budgeted": total, | |
| "selected_parts_complete": False, | |
| "complete_snapshot_scope": False, | |
| "columns": [], | |
| } | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| if resume and checkpoint.exists(): | |
| state = json.loads(checkpoint.read_text()) | |
| if state.get("signature") != signature: | |
| raise ValueError("Snapshot, source files or scope changed; use a new output/checkpoint") | |
| if not output.exists() or output.stat().st_size < state["bytes"]: | |
| raise ValueError("Checkpoint output is missing or truncated") | |
| elif output.exists(): | |
| raise FileExistsError("Snapshot output exists; use --resume or choose a new path") | |
| if not output.exists(): | |
| output.touch() | |
| if output.stat().st_size != state["bytes"]: | |
| with output.open("r+b") as handle: | |
| handle.truncate(state["bytes"]) | |
| atomic_json(checkpoint, state) | |
| if state["selected_parts_complete"] or (limit is not None and state["records"] >= limit): | |
| return state | |
| connection = _connection(output.parent / ".duckdb-extensions") | |
| try: | |
| if any(source.startswith("https://") for source in sources): | |
| # HTTP URLs avoid AWS signing and credential discovery altogether. | |
| connection.execute("INSTALL httpfs FROM core") | |
| connection.execute("LOAD httpfs") | |
| if work_ids is not None: | |
| connection.execute("CREATE TEMP TABLE wanted(id VARCHAR PRIMARY KEY)") | |
| connection.executemany( | |
| "INSERT INTO wanted VALUES (?)", [(x,) for x in sorted(work_ids)] | |
| ) | |
| while state["part"] < len(sources): | |
| if shutil.disk_usage(output.parent).free < min_free_bytes: | |
| raise RuntimeError( | |
| "Insufficient free space on the batch host; checkpoint preserved" | |
| ) | |
| source = sources[state["part"]] | |
| relation = connection.read_parquet(source) | |
| columns = set(relation.columns) | |
| if "id" not in columns or not {"title", "display_name"} & columns: | |
| raise ValueError( | |
| "Unexpected snapshot schema: id and title/display_name are required" | |
| ) | |
| projected = [field for field in PROJECTED_FIELDS if field in columns] | |
| state["columns"] = sorted(set(state["columns"]) | set(projected)) | |
| predicate = ( | |
| "regexp_replace(id, '^.*/', '') IN (SELECT id FROM wanted)" | |
| if work_ids is not None | |
| else coarse_predicate(columns, profile) | |
| ) | |
| # Native TIMESTAMPTZ conversion otherwise requires an unrelated | |
| # Python timezone package. ISO strings also serialize consistently. | |
| expressions = [ | |
| f'CAST("{field}" AS VARCHAR) AS "{field}"' | |
| if field in {"publication_date", "updated_date"} | |
| else f'"{field}"' | |
| for field in projected | |
| ] | |
| selected = relation.filter(predicate).project(", ".join(expressions)) | |
| if state["part_rows"]: | |
| selected = selected.limit(2**63 - 1, offset=state["part_rows"]) | |
| cursor = selected.execute() | |
| while rows := cursor.fetchmany(batch_size): | |
| if shutil.disk_usage(output.parent).free < min_free_bytes: | |
| raise RuntimeError( | |
| "Insufficient free space on the batch host; checkpoint preserved" | |
| ) | |
| reached_limit = False | |
| with output.open("ab") as handle: | |
| for row in rows: | |
| document = dict(zip(projected, row, strict=True)) | |
| state["part_rows"] += 1 | |
| if work_ids is None and not matches_work(document, profile, topic): | |
| continue | |
| document["snapshot_provenance"] = { | |
| "source": "public_s3", | |
| "manifest_sha256": snapshot_id, | |
| "snapshot_date": snapshot_date, | |
| "part": source, | |
| "scope": "reference_backfill" if work_ids is not None else profile, | |
| } | |
| handle.write( | |
| (json.dumps(document, ensure_ascii=False, default=str) + "\n").encode() | |
| ) | |
| state["records"] += 1 | |
| if not document.get("abstract_inverted_index"): | |
| state["metadata_only"] += 1 | |
| if limit is not None and state["records"] >= limit: | |
| reached_limit = True | |
| break | |
| handle.flush() | |
| os.fsync(handle.fileno()) | |
| state["bytes"] = handle.tell() | |
| atomic_json(checkpoint, state) | |
| if progress: | |
| progress(dict(state)) | |
| if reached_limit: | |
| return state # A row-limited run never claims source exhaustion. | |
| state.update(part=state["part"] + 1, part_rows=0) | |
| atomic_json(checkpoint, state) | |
| if progress: | |
| progress(dict(state)) | |
| state["selected_parts_complete"] = True | |
| state["complete_snapshot_scope"] = bool(all_parts_selected) | |
| atomic_json(checkpoint, state) | |
| return state | |
| finally: | |
| connection.close() | |