social-data-attribution-w2/tests/test_materialize_sample.py

import io
import json
from pathlib import Path

import pandas as pd
import pytest
import zstandard as zstd

from dolma.materialize_sample import (
    MaterializeResult,
    ShardStats,
    _iter_compressed_records,
    _output_path_for_shard,
    _slice_for_chunk,
    build_shard_index,
    materialize_all,
    materialize_shard,
    write_materialize_stats,
)
from dolma.writer import is_complete


def _make_jsonl_zst(records: list[dict]) -> bytes:
    cctx = zstd.ZstdCompressor(level=3)
    buf = io.BytesIO()
    with cctx.stream_writer(buf, closefd=False) as writer:
        for record in records:
            line = json.dumps(record) + "\n"
            writer.write(line.encode("utf-8"))
    return buf.getvalue()


@pytest.fixture
def sample_records():
    return [
        {"id": "doc_001", "text": "First document text.", "metadata": {"source": "a"}},
        {"id": "doc_002", "text": "Second document text.", "metadata": {"source": "b"}},
        {"id": "doc_003", "text": "Third document text.", "metadata": {"source": "c"}},
        {"id": "doc_004", "text": "Fourth document text.", "metadata": {"source": "d"}},
        {"id": "doc_005", "text": "Fifth document text.", "metadata": {"source": "e"}},
    ]


@pytest.fixture
def shard_data(sample_records):
    return _make_jsonl_zst(sample_records)


@pytest.fixture
def sample_manifest():
    return pd.DataFrame(
        {
            "doc_id": ["doc_001", "doc_002", "doc_003", "doc_010", "doc_011"],
            "shard_path": [
                "soc127/phase1_pool_shared/data/cc-001/shard_00000001.jsonl.zst",
                "soc127/phase1_pool_shared/data/cc-001/shard_00000001.jsonl.zst",
                "soc127/phase1_pool_shared/data/cc-002/shard_00000002.jsonl.zst",
                "soc127/phase2_nonpool_final/data/stack_edu/shard_00000001.jsonl.zst",
                "soc127/phase2_nonpool_final/data/stack_edu/shard_00000001.jsonl.zst",
            ],
            "token_count": [100, 200, 300, 150, 250],
        }
    )


class TestBuildShardIndex:
    def test_groups_by_shard_path(self, sample_manifest):
        index = build_shard_index(sample_manifest)
        assert len(index) == 3
        cc1_key = "soc127/phase1_pool_shared/data/cc-001/shard_00000001.jsonl.zst"
        assert index[cc1_key] == {"doc_001", "doc_002"}

    def test_missing_shard_path_column(self):
        df = pd.DataFrame({"doc_id": ["a"], "token_count": [100]})
        with pytest.raises(ValueError, match="shard_path"):
            build_shard_index(df)

    def test_missing_doc_id_column(self):
        df = pd.DataFrame({"shard_path": ["a"], "token_count": [100]})
        with pytest.raises(ValueError, match="doc_id"):
            build_shard_index(df)

    def test_empty_manifest(self):
        df = pd.DataFrame({"doc_id": [], "shard_path": []})
        index = build_shard_index(df)
        assert index == {}


class TestMaterializeShard:
    def test_extracts_matching_docs(self, shard_data, tmp_path):
        target_ids = {"doc_001", "doc_003", "doc_005"}
        out_path = tmp_path / "output.jsonl.zst"

        stats = materialize_shard(
            shard_data, "data/cc/shard.jsonl.zst", target_ids, out_path
        )

        assert stats.found_docs == 3
        assert stats.missing_docs == 0
        assert stats.expected_docs == 3
        assert stats.bytes_written > 0
        assert stats.missing_ids == []

    def test_reports_missing_docs(self, shard_data, tmp_path):
        target_ids = {"doc_001", "doc_999"}
        out_path = tmp_path / "output.jsonl.zst"

        stats = materialize_shard(
            shard_data, "data/cc/shard.jsonl.zst", target_ids, out_path
        )

        assert stats.found_docs == 1
        assert stats.missing_docs == 1
        assert stats.missing_ids == ["doc_999"]

    def test_creates_done_marker(self, shard_data, tmp_path):
        target_ids = {"doc_002"}
        out_path = tmp_path / "output.jsonl.zst"

        materialize_shard(shard_data, "data/cc/shard.jsonl.zst", target_ids, out_path)

        assert is_complete(out_path)

    def test_output_contains_correct_records(self, shard_data, tmp_path):
        target_ids = {"doc_002", "doc_004"}
        out_path = tmp_path / "output.jsonl.zst"

        materialize_shard(shard_data, "data/cc/shard.jsonl.zst", target_ids, out_path)

        dctx = zstd.ZstdDecompressor()
        with out_path.open("rb") as f:
            with dctx.stream_reader(f, read_across_frames=True) as reader:
                with io.TextIOWrapper(reader, encoding="utf-8") as text:
                    lines = [json.loads(line) for line in text if line.strip()]

        extracted_ids = {rec["id"] for rec in lines}
        assert extracted_ids == {"doc_002", "doc_004"}
        assert lines[0]["text"] == "Second document text."
        assert lines[1]["text"] == "Fourth document text."

    def test_no_matching_docs(self, shard_data, tmp_path):
        target_ids = {"doc_999", "doc_888"}
        out_path = tmp_path / "output.jsonl.zst"

        stats = materialize_shard(
            shard_data, "data/cc/shard.jsonl.zst", target_ids, out_path
        )

        assert stats.found_docs == 0
        assert stats.missing_docs == 2
        assert sorted(stats.missing_ids) == ["doc_888", "doc_999"]

    def test_all_docs_matched(self, sample_records, tmp_path):
        all_ids = {r["id"] for r in sample_records}
        shard_bytes = _make_jsonl_zst(sample_records)
        out_path = tmp_path / "output.jsonl.zst"

        stats = materialize_shard(
            shard_bytes, "data/cc/shard.jsonl.zst", all_ids, out_path
        )

        assert stats.found_docs == 5
        assert stats.missing_docs == 0


class TestSliceForChunk:
    def test_single_chunk(self):
        items = ["a", "b", "c", "d", "e"]
        assert _slice_for_chunk(items, 0, 1) == items

    def test_even_split(self):
        items = ["a", "b", "c", "d"]
        assert _slice_for_chunk(items, 0, 2) == ["a", "b"]
        assert _slice_for_chunk(items, 1, 2) == ["c", "d"]

    def test_uneven_split(self):
        items = ["a", "b", "c", "d", "e"]
        c0 = _slice_for_chunk(items, 0, 3)
        c1 = _slice_for_chunk(items, 1, 3)
        c2 = _slice_for_chunk(items, 2, 3)
        assert c0 + c1 + c2 == items
        assert len(c0) == 2
        assert len(c1) == 2
        assert len(c2) == 1

    def test_more_chunks_than_items(self):
        items = ["a", "b"]
        c0 = _slice_for_chunk(items, 0, 5)
        c1 = _slice_for_chunk(items, 1, 5)
        c4 = _slice_for_chunk(items, 4, 5)
        assert c0 == ["a"]
        assert c1 == ["b"]
        assert c4 == []

    def test_invalid_chunk_index(self):
        with pytest.raises(ValueError, match="out of range"):
            _slice_for_chunk(["a"], 1, 1)

    def test_invalid_chunk_count(self):
        with pytest.raises(ValueError, match="positive"):
            _slice_for_chunk(["a"], 0, 0)


class TestOutputPathForShard:
    def test_replaces_slashes(self):
        result = _output_path_for_shard(
            Path("/out"), "soc127/phase1/data/shard.jsonl.zst"
        )
        assert result.name == "soc127__phase1__data__shard.jsonl.zst"

    def test_adds_extension_if_missing(self):
        result = _output_path_for_shard(Path("/out"), "soc127/phase1/data/shard")
        assert result.name.endswith(".jsonl.zst")


class TestWriteMaterializeStats:
    def test_writes_stats_json(self, tmp_path):
        result = MaterializeResult(
            output_dir=tmp_path,
            total_expected=100,
            total_found=95,
            total_missing=5,
            total_bytes=1024,
            elapsed_seconds=10.5,
            missing_doc_ids=["id_1", "id_2"],
        )
        result.shard_stats = [
            ShardStats(shard_path="s1", expected_docs=50, found_docs=48),
            ShardStats(shard_path="s2", expected_docs=50, found_docs=47),
        ]

        stats_path = write_materialize_stats(result, tmp_path)

        assert stats_path.exists()
        payload = json.loads(stats_path.read_text())
        assert payload["total_expected_docs"] == 100
        assert payload["total_found_docs"] == 95
        assert payload["total_missing_docs"] == 5
        assert payload["shards_processed"] == 2
        assert payload["missing_doc_ids"] == ["id_1", "id_2"]


class TestIterCompressedRecords:
    def test_parses_valid_jsonl(self):
        records = [{"id": "a", "text": "hello"}, {"id": "b", "text": "world"}]
        data = _make_jsonl_zst(records)
        parsed = _iter_compressed_records(data)
        assert len(parsed) == 2
        assert parsed[0][1]["id"] == "a"
        assert parsed[1][1]["id"] == "b"

    def test_handles_malformed_json(self):
        cctx = zstd.ZstdCompressor(level=3)
        buf = io.BytesIO()
        with cctx.stream_writer(buf, closefd=False) as writer:
            writer.write(b'{"id": "a"}\n')
            writer.write(b"not valid json\n")
            writer.write(b'{"id": "c"}\n')
        data = buf.getvalue()
        parsed = _iter_compressed_records(data)
        assert len(parsed) == 3
        assert parsed[0][1]["id"] == "a"
        assert parsed[1][1] is None
        assert parsed[2][1]["id"] == "c"

    def test_empty_shard(self):
        data = _make_jsonl_zst([])
        parsed = _iter_compressed_records(data)
        assert parsed == []


class TestMaterializeAll:
    @pytest.fixture
    def two_shard_corpus(self, tmp_path):
        shard_a_records = [
            {"id": "doc_001", "text": "Alpha.", "metadata": {}},
            {"id": "doc_002", "text": "Beta.", "metadata": {}},
        ]
        shard_b_records = [
            {"id": "doc_003", "text": "Gamma.", "metadata": {}},
            {"id": "doc_004", "text": "Delta.", "metadata": {}},
        ]
        shard_a_path = "data/shard_a.jsonl.zst"
        shard_b_path = "data/shard_b.jsonl.zst"

        corpus_dir = tmp_path / "corpus"
        (corpus_dir / "data").mkdir(parents=True)
        (corpus_dir / shard_a_path).write_bytes(_make_jsonl_zst(shard_a_records))
        (corpus_dir / shard_b_path).write_bytes(_make_jsonl_zst(shard_b_records))

        manifest = pd.DataFrame(
            {
                "doc_id": ["doc_001", "doc_003", "doc_004"],
                "shard_path": [shard_a_path, shard_b_path, shard_b_path],
            }
        )
        manifest_path = tmp_path / "manifest.parquet"
        manifest.to_parquet(manifest_path, index=False)

        output_dir = tmp_path / "output"
        output_dir.mkdir()

        return {
            "corpus_dir": corpus_dir,
            "manifest_path": manifest_path,
            "output_dir": output_dir,
        }

    def test_materializes_docs_from_local_corpus(self, two_shard_corpus):
        result = materialize_all(
            manifest_path=two_shard_corpus["manifest_path"],
            output_dir=two_shard_corpus["output_dir"],
            r2_client=None,
            bucket="unused",
            corpus_dir=two_shard_corpus["corpus_dir"],
        )
        assert result.total_expected == 3
        assert result.total_found == 3
        assert result.total_missing == 0
        assert len(result.shard_stats) == 2

    def test_chunked_materialization_covers_all_shards(self, two_shard_corpus):
        results = []
        for chunk_idx in range(2):
            result = materialize_all(
                manifest_path=two_shard_corpus["manifest_path"],
                output_dir=two_shard_corpus["output_dir"],
                r2_client=None,
                bucket="unused",
                chunk_index=chunk_idx,
                chunk_count=2,
                corpus_dir=two_shard_corpus["corpus_dir"],
            )
            results.append(result)

        total_found = sum(r.total_found for r in results)
        total_expected = sum(r.total_expected for r in results)
        assert total_found == 3
        assert total_expected == 3

    def test_writes_output_files(self, two_shard_corpus):
        materialize_all(
            manifest_path=two_shard_corpus["manifest_path"],
            output_dir=two_shard_corpus["output_dir"],
            r2_client=None,
            bucket="unused",
            corpus_dir=two_shard_corpus["corpus_dir"],
        )
        output_files = list(two_shard_corpus["output_dir"].glob("*.jsonl.zst"))
        assert len(output_files) == 2
        for f in output_files:
            assert is_complete(f)