build_index.py

"""Build a metadata-only SQLite snapshot from an explicit, bounded source list."""
import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import sqlite3
import tempfile
import time

from huggingface_hub import HfApi, HfFileSystem, get_token
import pyarrow.parquet as pq

from catalog import ROOT, PROBS, normalize, unversioned, segment_metadata
from remote_catalog import BUCKET, MeasuredFile, source_info

SOURCES = [
    "annotations/fungi/s0000/s0000.chunk-00510.parquet",
    "annotations/protozoa/s0003/s0003.chunk-00017.parquet",
    "annotations/vertebrate_other/s0001/s0001.chunk-00018.parquet",
]


def build(output, sources):
    if not 1 <= len(sources) <= 10 or len(set(sources)) != len(sources):
        raise ValueError("Choose 1–10 distinct published Parquet objects.")
    if any(not s.startswith("annotations/") or not s.endswith(".parquet") for s in sources):
        raise ValueError("Only published annotation Parquet objects can be indexed.")
    api, fs = HfApi(token=get_token()), HfFileSystem(token=get_token())
    output = Path(output)
    output.parent.mkdir(parents=True, exist_ok=True)
    fd, name = tempfile.mkstemp(suffix=".sqlite", dir=output.parent)
    os.close(fd)
    temporary = Path(name)
    conn = sqlite3.connect(temporary)
    try:
        conn.executescript("""
            CREATE TABLE segments (
                id INTEGER PRIMARY KEY, assembly_accession TEXT NOT NULL,
                record_name TEXT NOT NULL, segment_start_bp INTEGER NOT NULL,
                segment_end_bp INTEGER NOT NULL, object_path TEXT NOT NULL,
                object_hash TEXT NOT NULL, row_group INTEGER NOT NULL,
                row_in_group INTEGER NOT NULL, metadata_json TEXT NOT NULL);
            CREATE TABLE aliases (alias TEXT NOT NULL, segment_id INTEGER NOT NULL,
                PRIMARY KEY (alias, segment_id)) WITHOUT ROWID;
            CREATE TABLE metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL);
        """)
        manifest = {"bucket_id": BUCKET, "created_at": datetime.now(timezone.utc).isoformat(),
                    "scope": "Explicit subset of published objects; not the full bucket or complete assemblies",
                    "sources": [], "rows": 0, "bases": 0, "assemblies": [], "divisions": []}
        assemblies, divisions = set(), set()
        for path in sources:
            source = source_info(api, BUCKET, path)
            if source is None:
                raise ValueError(f"Missing source: {path}")
            started = time.perf_counter()
            with MeasuredFile(fs, f"buckets/{BUCKET}/{path}") as remote:
                pf = pq.ParquetFile(remote)
                columns = [c for c in pf.schema_arrow.names if c not in PROBS and c != "sequence"]
                group_sizes = []
                for group in range(pf.num_row_groups):
                    group_sizes.append(pf.metadata.row_group(group).total_byte_size)
                    rows = pf.read_row_group(group, columns=columns, use_threads=False).to_pylist()
                    for row_in_group, record in enumerate(rows):
                        record = segment_metadata(record)
                        index = manifest["rows"]
                        conn.execute("INSERT INTO segments VALUES (?,?,?,?,?,?,?,?,?,?)", (
                            index, record["assembly_accession"], record["record_name"],
                            record["segment_start_bp"], record["segment_end_bp"], path,
                            source.xet_hash, group, row_in_group, json.dumps(record)))
                        aliases = {normalize(record["source_key"])}
                        for field in ("assembly_accession", "record_name"):
                            exact = normalize(record[field])
                            aliases.update([exact, unversioned(exact)])
                        conn.executemany("INSERT INTO aliases VALUES (?,?)", [(a, index) for a in aliases if a])
                        manifest["rows"] += 1
                        manifest["bases"] += record["segment_end_bp"] - record["segment_start_bp"]
                        assemblies.add(record["assembly_accession"])
                        divisions.add(record["division"])
                item = {"path": path, "xet_hash": source.xet_hash, "object_bytes": source.size,
                        "rows": pf.metadata.num_rows, "row_groups": pf.num_row_groups,
                        "max_uncompressed_row_group_bytes": max(group_sizes),
                        "metadata_bytes_fetched": remote.bytes_read, "range_reads": remote.range_reads,
                        "index_seconds": round(time.perf_counter() - started, 3)}
            after = source_info(api, BUCKET, path)
            if after is None or after.xet_hash != source.xet_hash:
                raise ValueError(f"Source changed while indexing: {path}")
            manifest["sources"].append(item)
            print(json.dumps(item), flush=True)
        manifest.update(assemblies=sorted(assemblies), divisions=sorted(divisions))
        conn.execute("INSERT INTO metadata VALUES ('manifest',?)", (json.dumps(manifest),))
        conn.commit()
        if conn.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
            raise ValueError("SQLite integrity check failed")
        conn.close()
        os.replace(temporary, output)
        print(json.dumps({"index_bytes": output.stat().st_size, "rows": manifest["rows"],
                          "assemblies": len(assemblies), "bases": manifest["bases"]}), flush=True)
        return manifest
    finally:
        conn.close()
        temporary.unlink(missing_ok=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", action="append", help="Repeat to select up to 10 published objects")
    parser.add_argument("--output", type=Path, default=ROOT / "data" / "catalog.sqlite")
    args = parser.parse_args()
    build(args.output, args.source or SOURCES)