social-data-attribution-w2/tests/quality/validation/test_io.py

from __future__ import annotations

import json

import pyarrow as pa
import pyarrow.parquet as pq
import pytest
import zstandard as zstd

from dolma.constants import FORMATS, TOPICS
from dolma.quality.sidecar import encode_quality_rows
from dolma.quality.validation.io import (
    read_quality_rows,
    read_raw_doc_map,
    read_soc91_doc_map,
)


def _encode_parquet_rows(rows: list[dict[str, object]]) -> bytes:
    sink = pa.BufferOutputStream()
    pq.write_table(pa.Table.from_pylist(rows), sink)
    return sink.getvalue().to_pybytes()


def _encode_zst_rows(rows: list[dict[str, object]]) -> bytes:
    raw = "".join(json.dumps(row) + "\n" for row in rows).encode("utf-8")
    return zstd.ZstdCompressor().compress(raw)


def test_read_quality_rows_round_trips_parquet_bytes(monkeypatch) -> None:
    payload = encode_quality_rows(
        [
            {
                "doc_id": "doc-1",
                "quality_label_id": 1,
                "quality_score": 0.9,
                "quality_high_prob": 0.9,
                "quality_low_prob": 0.1,
                "quality_confidence": 0.9,
            }
        ]
    )
    monkeypatch.setattr(
        "dolma.quality.validation.io.download_object_bytes",
        lambda *_args, **_kwargs: payload,
    )

    rows = read_quality_rows(
        object(),
        bucket="bucket",
        source_key="soc127/phase1_pool_shared/data/c4-en/shard_000.jsonl.zst",
        output_prefix="soc139-quality-validation/test",
    )

    assert len(rows) == 1
    row = rows[0]
    assert row["doc_id"] == "doc-1"
    assert row["quality_label_id"] == 1
    assert row["quality_score"] == pytest.approx(0.9)
    assert row["quality_high_prob"] == pytest.approx(0.9)
    assert row["quality_low_prob"] == pytest.approx(0.1)
    assert row["quality_confidence"] == pytest.approx(0.9)


def test_read_soc91_doc_map_maps_label_ids(monkeypatch) -> None:
    payload = _encode_parquet_rows(
        [
            {
                "doc_id": "doc-1",
                "topic_url_label_id": 0,
                "format_url_label_id": 16,
            }
        ]
    )
    monkeypatch.setattr(
        "dolma.quality.validation.io.download_object_bytes",
        lambda *_args, **_kwargs: payload,
    )

    joined = read_soc91_doc_map(
        object(),
        bucket="bucket",
        source_key="soc127/phase1_pool_shared/data/c4-en/shard_000.jsonl.zst",
        soc91_prefix="soc91-labels",
    )

    assert joined == {
        "doc-1": {
            "topic_url_label": TOPICS[0],
            "format_url_label": FORMATS[16],
        }
    }


def test_read_raw_doc_map_uses_soc127_doc_id_fallback(monkeypatch) -> None:
    payload = _encode_zst_rows(
        [
            {
                "_soc_127": {"doc_id": "mix-1", "source_family": "common_crawl"},
                "text": "hello world",
                "url": "https://example.com",
            }
        ]
    )
    monkeypatch.setattr(
        "dolma.quality.validation.io.download_object_bytes",
        lambda *_args, **_kwargs: payload,
    )

    rows = read_raw_doc_map(
        object(),
        bucket="bucket",
        source_key="soc127/phase2_nonpool_final/bucket_0000.jsonl.zst",
    )

    assert rows == {
        "mix-1": {
            "source_key": "soc127/phase2_nonpool_final/bucket_0000.jsonl.zst",
            "source_family": "common_crawl",
            "text_snippet": "hello world",
            "url": "https://example.com",
        }
    }