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.
70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
"""Dataset Export: snapshot features + labels to Parquet."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
|
|
from libs.common.config import get_settings
|
|
from libs.common.logging import configure_logging, get_logger
|
|
from libs.db.session import get_session
|
|
from libs.export.snapshot_export import export_dataset_snapshot
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
async def run_dataset_export(
|
|
snapshot_id: str | None,
|
|
split_policy: str,
|
|
output_dir: str,
|
|
) -> dict:
|
|
async with get_session() as session:
|
|
manifest = await export_dataset_snapshot(
|
|
session=session,
|
|
snapshot_id=snapshot_id,
|
|
split_policy=split_policy,
|
|
output_dir=output_dir,
|
|
)
|
|
return manifest
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Dataset Export")
|
|
parser.add_argument("--snapshot-id", default=None, help="Snapshot ID (UUID, auto-generated if omitted)")
|
|
parser.add_argument(
|
|
"--split-policy",
|
|
default="temporal_70_15_15",
|
|
help="Split policy string (default: temporal_70_15_15)",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
default="./data/datasets/snapshots",
|
|
help="Output directory for Parquet files",
|
|
)
|
|
parser.add_argument("--json", action="store_true", help="Print manifest JSON to stdout")
|
|
args = parser.parse_args()
|
|
|
|
settings = get_settings()
|
|
configure_logging(settings.log_level)
|
|
|
|
manifest = asyncio.run(
|
|
run_dataset_export(
|
|
snapshot_id=args.snapshot_id,
|
|
split_policy=args.split_policy,
|
|
output_dir=args.output_dir,
|
|
)
|
|
)
|
|
|
|
if args.json:
|
|
print(json.dumps(manifest, indent=2))
|
|
else:
|
|
print(f"Snapshot exported: {manifest['snapshot_id']}")
|
|
print(f" Output: {manifest['output_dir']}")
|
|
print(f" Total rows: {manifest['total_rows']}")
|
|
for split, count in manifest["row_counts"].items():
|
|
print(f" {split}: {count} rows")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|