tests/test_data_sharding.py

from __future__ import annotations

import subprocess
import sys
from pathlib import Path

from dovla_cil.data.schema import (
    CIL_VERSION,
    ActionChunk,
    CILRecord,
    RewardInfo,
    StructuredEffect,
    make_record_id,
)
from dovla_cil.data.sharding import (
    ShardReader,
    ShardWriter,
    iter_cil_records,
    split_records_by_group,
    write_cil_shards,
)
from dovla_cil.utils.io import iter_jsonl, read_json


def make_record(group_id: str, index: int, *, state_hash: str = "state-a") -> CILRecord:
    action = ActionChunk(
        action_id=f"action-{index}",
        representation="delta_xy",
        horizon=1,
        values=[[float(index), 0.0]],
        skill_type="push",
    )
    return CILRecord(
        version=CIL_VERSION,
        record_id=make_record_id(group_id, action.action_id, seed=0),
        group_id=group_id,
        state_hash=state_hash,
        task_id="task-a",
        scene_id=None,
        instruction="reach",
        instruction_family={"family": "toy"},
        observation_ref=None,
        observation_inline={"index": index},
        action_chunk=action,
        next_observation_ref=None,
        next_observation_inline={"index": index + 1},
        structured_effect=StructuredEffect(
            object_pose_delta={"object": [1.0, 0.0, 0.0]},
            moved_objects=["object"],
        ),
        reward=RewardInfo(
            progress=float(index),
            success=index > 0,
            terminal_success=index > 0,
            dense_components={"progress": float(index)},
        ),
        regret=None,
        rank_within_group=None,
        candidate_type="test",
        failure=None,
    )


def test_split_keeps_groups_together() -> None:
    records = [make_record("a", 0), make_record("b", 1), make_record("a", 2)]
    shards = split_records_by_group(records, max_records_per_shard=2)
    locations: dict[str, int] = {}
    for shard_index, shard in enumerate(shards):
        for record in shard:
            if record.group_id in locations:
                assert locations[record.group_id] == shard_index
            locations[record.group_id] = shard_index


def test_write_shards_roundtrip(tmp_path: Path) -> None:
    records = [make_record("a", 0), make_record("a", 1), make_record("b", 2)]
    manifest = write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=2)
    restored = []
    for shard in manifest["shards"]:
        restored.extend(iter_cil_records(tmp_path / str(shard["path"])))
    assert manifest["num_records"] == 3
    assert manifest["record_count"] == 3
    assert (tmp_path / "metadata.json").exists()
    assert (tmp_path / "manifest.json").exists()
    assert (tmp_path / "group_index.jsonl").exists()
    assert (tmp_path / "record_index.jsonl").exists()
    assert restored == records


def test_shard_writer_and_reader_roundtrip(tmp_path: Path) -> None:
    records = [
        make_record("group-a", 0),
        make_record("group-a", 1),
        make_record("group-b", 2, state_hash="state-b"),
    ]
    writer = ShardWriter(
        tmp_path,
        dataset_name="toy_dataset",
        backend="toy",
        k=2,
        task_count=1,
        seed=7,
        shard_size=2,
    )
    for record in records:
        writer.write(record)
    metadata = writer.close()

    reader = ShardReader(tmp_path)
    restored = list(reader.iterate_records())
    assert metadata["dataset_name"] == "toy_dataset"
    assert metadata["backend"] == "toy"
    assert metadata["num_records"] == 3
    assert metadata["num_groups"] == 2
    assert metadata["shards"][0]["path"] == "shards/shard_000000.jsonl"
    assert restored == records


def test_group_index_is_correct_and_group_loads(tmp_path: Path) -> None:
    records = [
        make_record("group-a", 0),
        make_record("group-a", 1),
        make_record("group-b", 2, state_hash="state-b"),
    ]
    write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=2)

    group_index = {row["group_id"]: row for row in iter_jsonl(tmp_path / "group_index.jsonl")}
    group_a = group_index["group-a"]
    assert group_a["shard_path"] == "shards/shard_000000.jsonl"
    assert group_a["record_ids"] == [records[0].record_id, records[1].record_id]
    assert group_a["task_id"] == "task-a"
    assert group_a["state_hash"] == "state-a"
    assert group_a["num_records"] == 2
    assert group_a["max_reward"] == 1.0
    assert group_a["success_count"] == 1
    assert group_a["candidate_type_counts"] == {"test": 2}

    reader = ShardReader(tmp_path / "metadata.json")
    assert reader.load_group("group-a") == records[:2]
    assert list(reader.iterate_groups())[1] == [records[2]]


def test_metadata_layout_and_record_index(tmp_path: Path) -> None:
    records = [make_record("group-a", 0), make_record("group-a", 1)]
    write_cil_shards(
        records,
        output_dir=tmp_path,
        max_records_per_shard=10,
        dataset_name="cil_toy",
        backend="toy",
        k=2,
        task_count=1,
        seed=3,
    )
    metadata = read_json(tmp_path / "metadata.json")
    record_index = list(iter_jsonl(tmp_path / "record_index.jsonl"))

    assert metadata["dataset_name"] == "cil_toy"
    assert metadata["version"]
    assert metadata["created_at"]
    assert metadata["backend"] == "toy"
    assert metadata["num_groups"] == 1
    assert metadata["num_records"] == 2
    assert metadata["k"] == 2
    assert metadata["task_count"] == 1
    assert metadata["seed"] == 3
    assert metadata["schema_version"] == CIL_VERSION
    assert record_index[0]["record_id"] == records[0].record_id
    assert record_index[0]["shard_path"] == "shards/shard_000000.jsonl"


def test_inspect_script_prints_summary_and_ranking(tmp_path: Path) -> None:
    records = [make_record("group-a", 0), make_record("group-a", 1)]
    write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=10)

    result = subprocess.run(
        [sys.executable, "scripts/inspect_shard.py", str(tmp_path), "--group-id", "group-a"],
        check=True,
        text=True,
        capture_output=True,
    )

    assert "num_records: 2" in result.stdout
    assert "sample_group: group-a" in result.stdout
    assert "record_id\tcandidate_type\treward.progress\tsuccess\tregret\trank\tfailure.type" in result.stdout