social-data-attribution-w2/scripts/analysis/build_trackstar_doc_bin_map.py

#!/usr/bin/env python3
# pyright: reportMissingImports=false
"""Map TrackStar positional document IDs to WebOrganizer bins."""

from __future__ import annotations

import argparse
import csv
import json
import re
from pathlib import Path

import pyarrow.parquet as pq

DOC_ID_RE = re.compile(r"^shard_(\d+):(\d+)$")


def load_doc_ids(path: Path) -> list[str]:
    with path.open() as fh:
        return [row["doc_id"] for row in csv.DictReader(fh)]


def shard_offsets(build_root: Path) -> dict[str, tuple[int, int]]:
    offsets = {}
    cursor = 0
    for shard in sorted(build_root.glob("shard_*")):
        info_path = shard / "info.json"
        if not info_path.exists():
            continue
        info = json.loads(info_path.read_text())
        count = int(info["num_grads"])
        offsets[shard.name] = (cursor, count)
        cursor += count
    if not offsets:
        raise ValueError(f"No shard info.json files under {build_root}")
    return offsets


def global_rows(doc_ids: list[str], offsets: dict[str, tuple[int, int]]) -> dict[str, int]:
    rows = {}
    for doc_id in doc_ids:
        match = DOC_ID_RE.match(doc_id)
        if match is None:
            raise ValueError(f"Unsupported TrackStar doc_id: {doc_id}")
        shard = f"shard_{int(match.group(1)):04d}"
        row_index = int(match.group(2))
        offset, count = offsets[shard]
        if row_index >= count:
            raise ValueError(f"{doc_id} row index exceeds {count}")
        rows[doc_id] = offset + row_index
    return rows


def write_map(rows: dict[str, int], manifest: Path, output: Path) -> None:
    table = pq.read_table(manifest, columns=["bin_topic", "bin_format"])
    topic = table["bin_topic"]
    fmt = table["bin_format"]
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", newline="") as fh:
        writer = csv.DictWriter(
            fh,
            fieldnames=["doc_id", "global_row", "bin_topic", "bin_format"],
        )
        writer.writeheader()
        for doc_id, row_index in rows.items():
            writer.writerow(
                {
                    "doc_id": doc_id,
                    "global_row": row_index,
                    "bin_topic": topic[row_index].as_py(),
                    "bin_format": fmt[row_index].as_py(),
                }
            )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--doc-ids-csv", type=Path, required=True)
    parser.add_argument("--build-root", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    doc_ids = load_doc_ids(args.doc_ids_csv)
    offsets = shard_offsets(args.build_root)
    rows = global_rows(doc_ids, offsets)
    write_map(rows, args.manifest, args.output)
    print(f"wrote {args.output} ({len(rows)} rows)")


if __name__ == "__main__":
    main()