You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""Merge multiple snapshot directories and re-split them temporally."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pyarrow.parquet as pq
|
|
|
|
from libs.common.time_utils import utc_now
|
|
from libs.export.snapshot_export import _rows_to_table, _temporal_split
|
|
|
|
|
|
def export_merged_snapshot(
|
|
*,
|
|
source_snapshot_dirs: list[str | Path],
|
|
output_dir: str | Path,
|
|
snapshot_id: str,
|
|
split_policy: str = "temporal_70_15_15",
|
|
) -> dict[str, Any]:
|
|
rows: list[dict[str, Any]] = []
|
|
source_ids: list[str] = []
|
|
|
|
for source_dir in source_snapshot_dirs:
|
|
source_path = Path(source_dir)
|
|
source_ids.append(source_path.name)
|
|
signal_origin = "continuation" if "cont_" in source_path.name else "event_day"
|
|
for split in ("train", "valid", "test"):
|
|
table = pq.read_table(source_path / f"{split}.parquet")
|
|
for row in table.to_pylist():
|
|
merged_row = dict(row)
|
|
merged_row.setdefault("signal_origin", signal_origin)
|
|
rows.append(merged_row)
|
|
|
|
splits = _temporal_split(rows, split_policy)
|
|
out_path = Path(output_dir) / snapshot_id
|
|
out_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
row_counts: dict[str, int] = {}
|
|
for split_name, split_rows in splits.items():
|
|
pq.write_table(_rows_to_table(split_rows), out_path / f"{split_name}.parquet")
|
|
row_counts[split_name] = len(split_rows)
|
|
|
|
manifest = {
|
|
"snapshot_id": snapshot_id,
|
|
"created_at_utc": utc_now().isoformat(),
|
|
"source_snapshot_ids": source_ids,
|
|
"transform": "merged_snapshot",
|
|
"split_policy": split_policy,
|
|
"row_counts": row_counts,
|
|
"total_rows": sum(row_counts.values()),
|
|
"output_dir": str(out_path),
|
|
}
|
|
(out_path / "manifest.json").write_text(json.dumps(manifest, indent=2))
|
|
return manifest
|