social-data-attribution-w2/src/dolma/materialize_sample.py

"""Materialize working sample documents from R2 source shards (SOC-134)."""

from __future__ import annotations

import io
import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path

import pandas as pd
import zstandard as zstd

from dolma.dedup.materialize import resolve_record_doc_id
from dolma.writer import JsonlWriter, is_complete

logger = logging.getLogger(__name__)


@dataclass
class ShardStats:
    shard_path: str
    expected_docs: int
    found_docs: int = 0
    missing_docs: int = 0
    missing_ids: list[str] = field(default_factory=list)
    bytes_written: int = 0
    elapsed_seconds: float = 0.0


@dataclass
class MaterializeResult:
    output_dir: Path
    shard_stats: list[ShardStats] = field(default_factory=list)
    total_expected: int = 0
    total_found: int = 0
    total_missing: int = 0
    total_bytes: int = 0
    elapsed_seconds: float = 0.0
    missing_doc_ids: list[str] = field(default_factory=list)


def build_shard_index(manifest_df: pd.DataFrame) -> dict[str, set[str]]:
    if "shard_path" not in manifest_df.columns:
        raise ValueError("Manifest missing required column: shard_path")
    if "doc_id" not in manifest_df.columns:
        raise ValueError("Manifest missing required column: doc_id")
    index: dict[str, set[str]] = {}
    for shard_path, doc_id in zip(
        manifest_df["shard_path"], manifest_df["doc_id"], strict=True
    ):
        index.setdefault(str(shard_path), set()).add(str(doc_id))
    return index


def _output_path_for_shard(output_dir: Path, shard_path: str) -> Path:
    safe_name = shard_path.replace("/", "__")
    if not safe_name.endswith(".jsonl.zst"):
        safe_name += ".jsonl.zst"
    return output_dir / safe_name


def _iter_compressed_records(
    data: bytes,
) -> list[tuple[str, dict[str, object] | None]]:
    dctx = zstd.ZstdDecompressor()
    records: list[tuple[str, dict[str, object] | None]] = []
    with io.BytesIO(data) as buf:
        with dctx.stream_reader(buf, read_across_frames=True) as reader:
            with io.TextIOWrapper(reader, encoding="utf-8") as text_reader:
                for raw_line in text_reader:
                    try:
                        records.append((raw_line, json.loads(raw_line)))
                    except json.JSONDecodeError:
                        records.append((raw_line, None))
    return records


def materialize_shard(
    shard_data: bytes,
    shard_path: str,
    target_doc_ids: set[str],
    output_path: Path,
) -> ShardStats:
    t0 = time.monotonic()
    stats = ShardStats(
        shard_path=shard_path,
        expected_docs=len(target_doc_ids),
    )

    remaining = set(target_doc_ids)
    with JsonlWriter(output_path) as writer:
        for _raw_line, record in _iter_compressed_records(shard_data):
            if record is None:
                continue
            doc_id, _field = resolve_record_doc_id(record, shard_path)
            if doc_id is None:
                continue
            if doc_id not in remaining:
                continue
            writer.write(record)
            remaining.discard(doc_id)
            stats.found_docs += 1
            if not remaining:
                break

        writer.write_stats()
        writer.write_done()

    stats.missing_docs = len(remaining)
    stats.missing_ids = sorted(remaining)
    stats.bytes_written = output_path.stat().st_size if output_path.exists() else 0
    stats.elapsed_seconds = time.monotonic() - t0
    return stats


def _read_shard_from_volume(
    shard_path: str,
    corpus_dir: Path,
) -> bytes | None:
    local_path = corpus_dir / shard_path
    if local_path.exists() and local_path.stat().st_size > 0:
        logger.info("Reading shard from volume: %s", shard_path)
        return local_path.read_bytes()
    return None


def materialize_shard_from_r2(
    r2_client: object,
    bucket: str,
    shard_path: str,
    target_doc_ids: set[str],
    output_path: Path,
    skip_existing: bool = True,
    corpus_dir: Path | None = None,
) -> ShardStats:
    if skip_existing and is_complete(output_path):
        logger.info("Skipping completed shard: %s", shard_path)
        return ShardStats(
            shard_path=shard_path,
            expected_docs=len(target_doc_ids),
            found_docs=len(target_doc_ids),
            bytes_written=output_path.stat().st_size,
        )

    if corpus_dir is not None:
        volume_data = _read_shard_from_volume(shard_path, corpus_dir)
        if volume_data is not None:
            return materialize_shard(
                volume_data, shard_path, target_doc_ids, output_path
            )

    from dolma.quality.r2 import download_object_bytes

    max_retries = 3
    for attempt in range(max_retries):
        try:
            logger.info(
                "Downloading shard from R2: %s (attempt %d)", shard_path, attempt + 1
            )
            shard_data = download_object_bytes(r2_client, bucket=bucket, key=shard_path)
            return materialize_shard(
                shard_data, shard_path, target_doc_ids, output_path
            )
        except Exception:
            if attempt == max_retries - 1:
                raise
            wait = (attempt + 1) * 5
            logger.warning(
                "Shard download failed, retrying in %ds: %s", wait, shard_path
            )
            import time as _time

            _time.sleep(wait)
    raise RuntimeError(f"Unreachable: all retries exhausted for {shard_path}")


def materialize_all(
    manifest_path: Path,
    output_dir: Path,
    r2_client: object,
    bucket: str,
    chunk_index: int = 0,
    chunk_count: int = 1,
    skip_existing: bool = True,
    corpus_dir: Path | None = None,
) -> MaterializeResult:
    t0 = time.monotonic()
    manifest_df = pd.read_parquet(manifest_path)
    shard_index = build_shard_index(manifest_df)
    all_shards = sorted(shard_index.keys())

    if chunk_count > 1:
        chunk_shards = _slice_for_chunk(all_shards, chunk_index, chunk_count)
    else:
        chunk_shards = all_shards

    logger.info(
        "Materializing chunk %d/%d: %d shards (of %d total)",
        chunk_index,
        chunk_count,
        len(chunk_shards),
        len(all_shards),
    )

    output_dir.mkdir(parents=True, exist_ok=True)
    result = MaterializeResult(output_dir=output_dir)

    for i, shard_path in enumerate(chunk_shards):
        target_ids = shard_index[shard_path]
        out_path = _output_path_for_shard(output_dir, shard_path)

        stats = materialize_shard_from_r2(
            r2_client,
            bucket,
            shard_path,
            target_ids,
            out_path,
            skip_existing=skip_existing,
            corpus_dir=corpus_dir,
        )

        result.shard_stats.append(stats)
        result.total_expected += stats.expected_docs
        result.total_found += stats.found_docs
        result.total_missing += stats.missing_docs
        result.total_bytes += stats.bytes_written
        result.missing_doc_ids.extend(stats.missing_ids)

        if (i + 1) % 10 == 0 or i == len(chunk_shards) - 1:
            logger.info(
                "Progress: %d/%d shards, %d/%d docs found",
                i + 1,
                len(chunk_shards),
                result.total_found,
                result.total_expected,
            )

    result.elapsed_seconds = time.monotonic() - t0

    logger.info(
        "Materialization complete: %d/%d docs found across %d shards in %.1fs",
        result.total_found,
        result.total_expected,
        len(chunk_shards),
        result.elapsed_seconds,
    )

    return result


def _slice_for_chunk(items: list[str], chunk_index: int, chunk_count: int) -> list[str]:
    if chunk_count <= 0:
        raise ValueError("chunk_count must be positive")
    if chunk_index < 0 or chunk_index >= chunk_count:
        raise ValueError(f"chunk_index {chunk_index} out of range [0, {chunk_count})")
    chunk_size = len(items) // chunk_count
    remainder = len(items) % chunk_count
    start = chunk_index * chunk_size + min(chunk_index, remainder)
    end = start + chunk_size + (1 if chunk_index < remainder else 0)
    return items[start:end]


def write_materialize_stats(result: MaterializeResult, output_dir: Path) -> Path:
    stats_path = output_dir / "materialize_stats.json"
    payload = {
        "total_expected_docs": result.total_expected,
        "total_found_docs": result.total_found,
        "total_missing_docs": result.total_missing,
        "total_bytes_written": result.total_bytes,
        "elapsed_seconds": round(result.elapsed_seconds, 1),
        "shards_processed": len(result.shard_stats),
        "missing_doc_ids": result.missing_doc_ids[:100],
    }
    stats_path.parent.mkdir(parents=True, exist_ok=True)
    stats_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    return stats_path


__all__ = [
    "MaterializeResult",
    "ShardStats",
    "build_shard_index",
    "materialize_all",
    "materialize_shard",
    "materialize_shard_from_r2",
    "write_materialize_stats",
]