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

import json

import pandas as pd
import pytest

from dolma.paper.sampling_loader import apply_token_filter, load_exclusion_doc_ids
from dolma.pool_sample.sampling import generate_dummy_manifest
from dolma.working_sample import (
    NUM_BINS,
    BinResult,
    bin_summary_dataframe,
    draw_working_sample,
    sample_contract,
    write_outputs,
)


@pytest.fixture
def dummy_manifest():
    return generate_dummy_manifest(n_docs=10_000, seed=42)


def test_draw_working_sample_returns_expected_types(dummy_manifest):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    assert isinstance(result.sample_df, pd.DataFrame)
    assert len(result.bin_results) == NUM_BINS
    assert all(isinstance(br, BinResult) for br in result.bin_results)


def test_draw_working_sample_respects_token_floor(dummy_manifest):
    floor = 10_000
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=floor, seed=42)
    for br in result.bin_results:
        if br.available_tokens >= floor:
            assert br.realized_tokens >= floor


def test_draw_working_sample_exhausts_small_bins(dummy_manifest):
    result = draw_working_sample(
        dummy_manifest, token_floor_per_bin=10_000_000, seed=42
    )
    for br in result.bin_results:
        if br.available_docs > 0 and br.available_tokens < 10_000_000:
            assert br.realized_docs == br.available_docs
            assert br.underfilled is True


def test_draw_working_sample_deterministic(dummy_manifest):
    r1 = draw_working_sample(dummy_manifest, token_floor_per_bin=1000, seed=99)
    r2 = draw_working_sample(dummy_manifest, token_floor_per_bin=1000, seed=99)
    assert len(r1.sample_df) == len(r2.sample_df)
    assert set(r1.sample_df["doc_id"]) == set(r2.sample_df["doc_id"])


def test_bin_summary_shape(dummy_manifest):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    summary = bin_summary_dataframe(result)
    assert len(summary) == NUM_BINS
    expected_cols = {
        "bin_id",
        "topic",
        "format",
        "requested_floor_tokens",
        "requested_docs_per_bin",
        "realized_tokens",
        "realized_docs",
        "available_tokens",
        "available_docs",
        "underfilled",
        "shortfall_tokens",
        "shortfall_docs",
    }
    assert set(summary.columns) == expected_cols


def test_sample_contract_keys(dummy_manifest):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    contract = sample_contract(result)
    expected_keys = {
        "WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN",
        "WORKING_SAMPLE_DOCS_PER_BIN",
        "WORKING_SAMPLE_GLOBAL_TOKEN_BUDGET",
        "WORKING_SAMPLE_MIN_TOKEN_COUNT",
        "WORKING_SAMPLE_MAX_TOKEN_COUNT",
        "WORKING_SAMPLE_REALIZED_TOKEN_TOTAL",
        "WORKING_SAMPLE_REALIZED_DOC_COUNT",
        "WORKING_SAMPLE_UNDERFILLED_BIN_COUNT",
        "WORKING_SAMPLE_COVERED_BIN_COUNT",
        "WORKING_SAMPLE_TOTAL_BIN_COUNT",
        "WORKING_SAMPLE_SAMPLING_SEED",
    }
    assert set(contract.keys()) == expected_keys


def test_write_outputs_creates_files(dummy_manifest, tmp_path):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    write_outputs(result, tmp_path)
    assert (tmp_path / "working_sample_manifest.parquet").exists()
    assert (tmp_path / "bin_summary.csv").exists()
    assert (tmp_path / "sample_contract.json").exists()

    contract = json.loads((tmp_path / "sample_contract.json").read_text())
    assert contract["WORKING_SAMPLE_SAMPLING_SEED"] == 42


def test_sample_doc_ids_subset_of_manifest(dummy_manifest):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    manifest_ids = set(dummy_manifest["doc_id"])
    sample_ids = set(result.sample_df["doc_id"])
    assert sample_ids.issubset(manifest_ids)


def test_no_duplicate_doc_ids_in_sample(dummy_manifest):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=100, seed=42)
    assert result.sample_df["doc_id"].is_unique


@pytest.mark.parametrize(
    "floor,expected_underfilled_range",
    [
        (1, (0, 400)),
        (10_000_000, (400, 576)),
    ],
)
def test_underfilled_count_scales_with_floor(
    dummy_manifest, floor, expected_underfilled_range
):
    result = draw_working_sample(dummy_manifest, token_floor_per_bin=floor, seed=42)
    contract = sample_contract(result)
    lo, hi = expected_underfilled_range
    assert lo <= contract["WORKING_SAMPLE_UNDERFILLED_BIN_COUNT"] <= hi


def test_docs_per_bin_returns_expected_count(dummy_manifest):
    result = draw_working_sample(dummy_manifest, docs_per_bin=5, seed=42)
    for br in result.bin_results:
        if br.available_docs >= 5:
            assert br.realized_docs == 5
        elif br.available_docs > 0:
            assert br.realized_docs == br.available_docs


def test_docs_per_bin_deterministic(dummy_manifest):
    r1 = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=99)
    r2 = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=99)
    assert len(r1.sample_df) == len(r2.sample_df)
    assert set(r1.sample_df["doc_id"]) == set(r2.sample_df["doc_id"])


def test_docs_per_bin_caps_at_available(dummy_manifest):
    result = draw_working_sample(dummy_manifest, docs_per_bin=100_000, seed=42)
    for br in result.bin_results:
        if br.available_docs > 0:
            assert br.realized_docs == br.available_docs
            assert br.underfilled is True


def test_docs_per_bin_contract(dummy_manifest):
    result = draw_working_sample(dummy_manifest, docs_per_bin=5, seed=42)
    contract = sample_contract(result)
    assert contract["WORKING_SAMPLE_DOCS_PER_BIN"] == 5
    assert contract["WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN"] == 0


def test_docs_per_bin_no_duplicates(dummy_manifest):
    result = draw_working_sample(dummy_manifest, docs_per_bin=10, seed=42)
    assert result.sample_df["doc_id"].is_unique


def test_apply_token_filter_min(dummy_manifest):
    min_tokens = 512
    filtered = apply_token_filter(dummy_manifest, "token_count", min_tokens=min_tokens)
    assert len(filtered) < len(dummy_manifest)
    assert (filtered["token_count"] >= min_tokens).all()


def test_apply_token_filter_max(dummy_manifest):
    max_tokens = 2000
    filtered = apply_token_filter(dummy_manifest, "token_count", max_tokens=max_tokens)
    assert len(filtered) < len(dummy_manifest)
    assert (filtered["token_count"] <= max_tokens).all()


def test_apply_token_filter_noop(dummy_manifest):
    filtered = apply_token_filter(dummy_manifest, "token_count")
    assert len(filtered) == len(dummy_manifest)


def test_sample_contract_records_filter_values(dummy_manifest):
    filtered = apply_token_filter(dummy_manifest, "token_count", min_tokens=512)
    result = draw_working_sample(
        filtered,
        token_floor_per_bin=100,
        seed=42,
        min_token_count=512,
        max_token_count=None,
    )
    contract = sample_contract(result)
    assert contract["WORKING_SAMPLE_MIN_TOKEN_COUNT"] == 512
    assert contract["WORKING_SAMPLE_MAX_TOKEN_COUNT"] is None


class TestExclusionManifests:
    def test_load_exclusion_doc_ids_empty(self):
        assert load_exclusion_doc_ids([]) == set()

    def test_load_exclusion_doc_ids_single(self, tmp_path):
        manifest = pd.DataFrame({"doc_id": ["a", "b", "c"]})
        path = tmp_path / "exclude.parquet"
        manifest.to_parquet(path, index=False)
        ids = load_exclusion_doc_ids([path])
        assert ids == {"a", "b", "c"}

    def test_load_exclusion_doc_ids_multiple(self, tmp_path):
        m1 = pd.DataFrame({"doc_id": ["a", "b"]})
        m2 = pd.DataFrame({"doc_id": ["b", "c", "d"]})
        p1 = tmp_path / "e1.parquet"
        p2 = tmp_path / "e2.parquet"
        m1.to_parquet(p1, index=False)
        m2.to_parquet(p2, index=False)
        ids = load_exclusion_doc_ids([p1, p2])
        assert ids == {"a", "b", "c", "d"}

    def test_load_exclusion_doc_ids_missing_file(self, tmp_path):
        with pytest.raises(FileNotFoundError):
            load_exclusion_doc_ids([tmp_path / "nonexistent.parquet"])

    def test_exclusion_removes_docs_from_sample(self, dummy_manifest):
        result_a = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=42)
        excluded_ids = set(result_a.sample_df["doc_id"])

        remaining = dummy_manifest.loc[
            ~dummy_manifest["doc_id"].isin(excluded_ids)
        ].reset_index(drop=True)
        result_b = draw_working_sample(remaining, docs_per_bin=3, seed=42)

        overlap = excluded_ids & set(result_b.sample_df["doc_id"])
        assert len(overlap) == 0

    def test_staggered_samples_are_disjoint(self, dummy_manifest):
        result_a = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=42)
        ids_a = set(result_a.sample_df["doc_id"])

        remaining = dummy_manifest.loc[
            ~dummy_manifest["doc_id"].isin(ids_a)
        ].reset_index(drop=True)
        result_b = draw_working_sample(remaining, docs_per_bin=3, seed=42)
        ids_b = set(result_b.sample_df["doc_id"])

        assert ids_a.isdisjoint(ids_b)
        assert len(ids_a) > 0
        assert len(ids_b) > 0

    def test_exclusion_contract_fields(self, dummy_manifest):
        result = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=42)
        result.exclude_manifest_paths = ["/some/path.parquet"]
        result.excluded_doc_count = 100

        contract = sample_contract(result)
        assert contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] == [
            "/some/path.parquet"
        ]
        assert contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] == 100

    def test_no_exclusion_contract_fields_when_empty(self, dummy_manifest):
        result = draw_working_sample(dummy_manifest, docs_per_bin=3, seed=42)
        contract = sample_contract(result)
        assert "WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS" not in contract
        assert "WORKING_SAMPLE_EXCLUDED_DOC_COUNT" not in contract

    def test_exclusion_with_unknown_ids_is_graceful(self, dummy_manifest):
        bogus_ids = {"nonexistent_1", "nonexistent_2"}
        remaining = dummy_manifest.loc[
            ~dummy_manifest["doc_id"].isin(bogus_ids)
        ].reset_index(drop=True)
        result = draw_working_sample(remaining, docs_per_bin=3, seed=42)
        assert len(result.sample_df) > 0