from __future__ import annotations from collections import Counter, OrderedDict from datetime import datetime, timezone from pathlib import Path from typing import Iterable, Iterator from dovla_cil.data.schema import CILRecord, CIL_VERSION from dovla_cil.utils.io import ensure_dir, iter_jsonl, write_json, write_jsonl def group_records(records: Iterable[CILRecord]) -> OrderedDict[str, list[CILRecord]]: groups: OrderedDict[str, list[CILRecord]] = OrderedDict() for record in records: groups.setdefault(record.group_id, []).append(record) return groups def split_records_by_group( records: Iterable[CILRecord], *, max_records_per_shard: int ) -> list[list[CILRecord]]: if max_records_per_shard <= 0: raise ValueError("max_records_per_shard must be positive") shards: list[list[CILRecord]] = [] current: list[CILRecord] = [] for group in group_records(records).values(): if current and len(current) + len(group) > max_records_per_shard: shards.append(current) current = [] current.extend(group) if current: shards.append(current) return shards class ShardWriter: """Streaming writer for grouped CIL JSONL shards. Records should be written group-contiguously. This keeps a single CIL group in one shard and makes the group index compact and deterministic. If a closed group appears again, the writer raises an error instead of silently creating a broken index. """ def __init__( self, output_dir: str | Path, *, dataset_name: str = "dovla_cil", backend: str = "unknown", k: int | None = None, task_count: int | None = None, seed: int | None = None, shard_size: int = 1024, schema_version: str = CIL_VERSION, version: str = "0.1", shard_format: str = "jsonl", index_format: str = "jsonl", overwrite: bool = True, ) -> None: if shard_size <= 0: raise ValueError("shard_size must be positive") self.output_dir = ensure_dir(output_dir) self.shards_dir = ensure_dir(self.output_dir / "shards") self.dataset_name = dataset_name self.backend = backend self.k = k self.task_count = task_count self.seed = seed self.shard_size = int(shard_size) self.schema_version = schema_version self.version = version self.shard_format = shard_format.lower() self.index_format = index_format.lower() if self.shard_format not in {"jsonl", "parquet"}: raise ValueError("shard_format must be 'jsonl' or 'parquet'") if self.index_format not in {"jsonl", "parquet"}: raise ValueError("index_format must be 'jsonl' or 'parquet'") if overwrite: self._remove_stale_outputs() self._current_group_id: str | None = None self._current_group: list[CILRecord] = [] self._closed_groups: set[str] = set() self._current_shard: list[CILRecord] = [] self._current_shard_group_ids: list[str] = [] self._shard_index = 0 self._shard_entries: list[dict[str, object]] = [] self._group_index_entries: list[dict[str, object]] = [] self._record_index_entries: list[dict[str, object]] = [] self._task_ids: set[str] = set() self._num_records = 0 self._closed = False def write(self, record: CILRecord) -> None: if self._closed: raise RuntimeError("Cannot write to a closed ShardWriter") record.validate() self._task_ids.add(record.task_id) if self._current_group_id is None: self._current_group_id = record.group_id if record.group_id != self._current_group_id: self._flush_group() if record.group_id in self._closed_groups: raise ValueError( f"Group {record.group_id!r} was already written. " "ShardWriter requires group-contiguous records." ) self._current_group_id = record.group_id self._current_group.append(record) def close(self) -> dict[str, object]: if self._closed: return self._metadata() self._flush_group() self._flush_shard() metadata = self._metadata() self._write_indices_and_metadata(metadata) self._closed = True return metadata def _flush_group(self) -> None: if not self._current_group: return group = list(self._current_group) group_id = group[0].group_id if any(record.group_id != group_id for record in group): raise ValueError("Internal writer error: mixed group buffer") if self._current_shard and len(self._current_shard) + len(group) > self.shard_size: self._flush_shard() shard_path = self._current_shard_path() row_offset = len(self._current_shard) self._current_shard.extend(group) self._current_shard_group_ids.append(group_id) self._group_index_entries.append(_make_group_index_entry(group, shard_path=shard_path)) for row_index, record in enumerate(group, start=row_offset): self._record_index_entries.append( _make_record_index_entry(record, shard_path=shard_path, row_index=row_index) ) self._num_records += len(group) self._closed_groups.add(group_id) self._current_group = [] self._current_group_id = None if len(self._current_shard) >= self.shard_size: self._flush_shard() def _flush_shard(self) -> None: if not self._current_shard: return shard_path = self._current_shard_path() target = self.output_dir / shard_path if self.shard_format == "jsonl": write_jsonl((record.to_dict() for record in self._current_shard), target) else: _write_parquet_rows([record.to_dict() for record in self._current_shard], target) self._shard_entries.append( { "path": shard_path, "record_count": len(self._current_shard), "group_ids": list(self._current_shard_group_ids), "format": self.shard_format, } ) self._current_shard = [] self._current_shard_group_ids = [] self._shard_index += 1 def _current_shard_path(self) -> str: suffix = "jsonl" if self.shard_format == "jsonl" else "parquet" return f"shards/shard_{self._shard_index:06d}.{suffix}" def _metadata(self) -> dict[str, object]: inferred_k = self.k if inferred_k is None and self._group_index_entries: inferred_k = max(int(entry["num_records"]) for entry in self._group_index_entries) if inferred_k is None: inferred_k = 0 task_count = self.task_count if self.task_count is not None else len(self._task_ids) return { "dataset_name": self.dataset_name, "version": self.version, "created_at": datetime.now(timezone.utc).isoformat(), "backend": self.backend, "num_groups": len(self._group_index_entries), "num_records": self._num_records, "k": int(inferred_k), "task_count": int(task_count), "seed": self.seed, "schema_version": self.schema_version, "format": "dovla_cil", "shard_format": self.shard_format, "index_format": self.index_format, "shard_size": self.shard_size, "shard_count": len(self._shard_entries), "group_index_path": _index_path("group_index", self.index_format), "record_index_path": _index_path("record_index", self.index_format), "shards": list(self._shard_entries), # Compatibility aliases for the original manifest shape. "record_count": self._num_records, "group_count": len(self._group_index_entries), } def _write_indices_and_metadata(self, metadata: dict[str, object]) -> None: _write_index_rows( self._group_index_entries, self.output_dir / str(metadata["group_index_path"]), index_format=self.index_format, ) _write_index_rows( self._record_index_entries, self.output_dir / str(metadata["record_index_path"]), index_format=self.index_format, ) write_json(metadata, self.output_dir / "metadata.json") write_json(metadata, self.output_dir / "manifest.json") def _remove_stale_outputs(self) -> None: for pattern in ( "metadata.json", "manifest.json", "group_index.jsonl", "record_index.jsonl", "group_index.parquet", "record_index.parquet", ): target = self.output_dir / pattern if target.exists() and target.is_file(): target.unlink() for target in self.shards_dir.glob("shard_*.*"): if target.is_file(): target.unlink() class ShardReader: """Read CIL datasets written by `ShardWriter`.""" def __init__(self, dataset_path: str | Path) -> None: from dovla_cil.data.index import ShardIndex self.index = ShardIndex.from_path(dataset_path) def iterate_records(self) -> Iterator[CILRecord]: for shard_path in self.index.shards: yield from iter_cil_records(shard_path) def iterate_groups(self) -> Iterator[list[CILRecord]]: for entry in self.index.load_group_index(): yield self.load_group(str(entry["group_id"])) def load_group(self, group_id: str) -> list[CILRecord]: entry = self.index.group_entry(group_id) shard_path = self.index.dataset_dir / str(entry["shard_path"]) records = [record for record in iter_cil_records(shard_path) if record.group_id == group_id] expected = int(entry["num_records"]) if len(records) != expected: raise ValueError( f"Group {group_id!r} index expected {expected} records, loaded {len(records)}" ) return records def write_cil_shards( records: Iterable[CILRecord], *, output_dir: str | Path, max_records_per_shard: int = 1024, prefix: str = "cil", dataset_name: str = "dovla_cil", backend: str = "unknown", k: int | None = None, task_count: int | None = None, seed: int | None = None, ) -> dict[str, object]: del prefix # ShardWriter uses the canonical shard_000000 layout. materialized = list(records) writer = ShardWriter( output_dir, dataset_name=dataset_name, backend=backend, k=k, task_count=task_count, seed=seed, shard_size=max_records_per_shard, ) for group in group_records(materialized).values(): for record in group: writer.write(record) return writer.close() def iter_cil_records(path: str | Path) -> Iterator[CILRecord]: record_path = Path(path) if record_path.is_dir() or record_path.name in {"metadata.json", "manifest.json"}: yield from ShardReader(record_path).iterate_records() return if record_path.suffix == ".jsonl": for payload in iter_jsonl(record_path): yield CILRecord.from_dict(payload) return if record_path.suffix == ".parquet": for payload in _read_parquet_rows(record_path): yield CILRecord.from_dict(payload) return raise ValueError(f"Unsupported CIL record path: {record_path}") def _make_group_index_entry(group: list[CILRecord], *, shard_path: str) -> dict[str, object]: first = group[0] candidate_counts = Counter(record.candidate_type for record in group) rewards = [record.reward.progress for record in group] return { "group_id": first.group_id, "shard_path": shard_path, "record_ids": [record.record_id for record in group], "task_id": first.task_id, "state_hash": first.state_hash, "num_records": len(group), "max_reward": max(rewards) if rewards else 0.0, "success_count": sum(1 for record in group if record.reward.terminal_success), "candidate_type_counts": dict(sorted(candidate_counts.items())), # Useful extras for inspection and compatibility with older group_index.jsonl files. "scene_id": first.scene_id, "instruction": first.instruction, "state_blob_ref": first.metadata.get("state_blob_ref"), } def _make_record_index_entry( record: CILRecord, *, shard_path: str, row_index: int ) -> dict[str, object]: return { "record_id": record.record_id, "group_id": record.group_id, "shard_path": shard_path, "row_index": row_index, "task_id": record.task_id, "state_hash": record.state_hash, "candidate_type": record.candidate_type, "reward_progress": record.reward.progress, "success": record.reward.terminal_success, "regret": record.regret, "rank_within_group": record.rank_within_group, "failure_type": record.failure.type if record.failure else None, } def _index_path(stem: str, index_format: str) -> str: suffix = "jsonl" if index_format == "jsonl" else "parquet" return f"{stem}.{suffix}" def _write_index_rows(rows: list[dict[str, object]], path: Path, *, index_format: str) -> None: if index_format == "jsonl": write_jsonl(rows, path) return _write_parquet_rows(rows, path) def _write_parquet_rows(rows: list[dict[str, object]], path: Path) -> None: try: import pandas as pd except ImportError as exc: # pragma: no cover - optional dependency path raise ImportError("Parquet output requires pandas and pyarrow.") from exc ensure_dir(path.parent) pd.DataFrame(rows).to_parquet(path, index=False) def _read_parquet_rows(path: Path) -> list[dict[str, object]]: try: import pandas as pd except ImportError as exc: # pragma: no cover - optional dependency path raise ImportError("Parquet input requires pandas and pyarrow.") from exc return pd.read_parquet(path).to_dict(orient="records")