| """Convert raw CSV/Parquet into qlib-compatible per-symbol CSV, then dump to .bin format.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| from config.settings import load_settings |
|
|
|
|
| REQUIRED_QLIB_FIELDS = ["open", "close", "high", "low", "volume", "factor"] |
|
|
|
|
| def normalize_columns(df: pd.DataFrame, aliases: dict[str, str]) -> pd.DataFrame: |
| rename = {} |
| for col in df.columns: |
| key = str(col).strip() |
| if key in aliases: |
| rename[col] = aliases[key] |
| out = df.rename(columns=rename) |
| return out.loc[:, ~out.columns.duplicated()] |
|
|
|
|
| def parquet_to_symbol_csv( |
| parquet_path: Path, |
| output_dir: Path, |
| aliases: dict[str, str], |
| chunk_size: int = 500_000, |
| ) -> int: |
| """Split a wide parquet into per-symbol CSV files for dump_bin.""" |
| import pyarrow.parquet as pq |
|
|
| output_dir.mkdir(parents=True, exist_ok=True) |
| pf = pq.ParquetFile(parquet_path) |
| writers: dict[str, list[pd.DataFrame]] = {} |
|
|
| for batch in pf.iter_batches(batch_size=chunk_size): |
| chunk = batch.to_pandas() |
| chunk = normalize_columns(chunk, aliases) |
| if "date" not in chunk.columns or "symbol" not in chunk.columns: |
| raise KeyError(f"Missing date/symbol after alias map. columns={chunk.columns.tolist()}") |
|
|
| chunk["date"] = pd.to_datetime(chunk["date"]) |
| chunk["symbol"] = chunk["symbol"].astype(str).str.lower() |
| if "factor" not in chunk.columns: |
| chunk["factor"] = 1.0 |
|
|
| for sym, grp in chunk.groupby("symbol"): |
| writers.setdefault(sym, []).append(grp) |
|
|
| count = 0 |
| for sym, parts in writers.items(): |
| df = pd.concat(parts, ignore_index=True) |
| df = df.sort_values("date").drop_duplicates("date") |
| out_path = output_dir / f"{sym}.csv" |
| keep = ["date", "symbol"] + [c for c in REQUIRED_QLIB_FIELDS if c in df.columns] |
| df[keep].to_csv(out_path, index=False) |
| count += 1 |
| return count |
|
|
|
|
| def dump_csv_dir_to_qlib( |
| csv_dir: Path, |
| qlib_dir: Path, |
| include_fields: str, |
| freq: str = "day", |
| max_workers: int = 8, |
| mode: str = "dump_all", |
| ): |
| from scripts.dump_bin import DumpDataAll, DumpDataUpdate |
|
|
| kwargs = dict( |
| data_path=str(csv_dir), |
| qlib_dir=str(qlib_dir), |
| include_fields=include_fields, |
| file_suffix=".csv", |
| date_field_name="date", |
| symbol_field_name="symbol", |
| freq=freq, |
| max_workers=max_workers, |
| ) |
| if mode == "dump_update": |
| DumpDataUpdate(**kwargs)() |
| else: |
| DumpDataAll(**kwargs)() |
|
|
|
|
| def run_from_parquet(parquet_path: Path | None = None, mode: str = "dump_all") -> Path: |
| settings = load_settings() |
| dump_cfg = settings.dump_config |
|
|
| csv_dir = settings.path(dump_cfg.get("raw_data_dir", "data/raw/csv_by_symbol")) |
| qlib_dir = settings.path(dump_cfg.get("qlib_dir", "data/qlib_data/cn_data")) |
| aliases = dump_cfg.get("column_aliases", {}) |
|
|
| if parquet_path is not None: |
| n = parquet_to_symbol_csv(Path(parquet_path), csv_dir, aliases) |
| print(f"Converted {n} symbols to {csv_dir}") |
|
|
| dump_csv_dir_to_qlib( |
| csv_dir=csv_dir, |
| qlib_dir=qlib_dir, |
| include_fields=dump_cfg.get("include_fields", "open,close,high,low,volume,factor"), |
| freq=dump_cfg.get("freq", "day"), |
| max_workers=int(dump_cfg.get("max_workers", 8)), |
| mode=mode, |
| ) |
| print(f"Qlib bin data ready at: {qlib_dir}") |
| return qlib_dir |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
|
|
| root = Path(__file__).resolve().parents[1] |
| if str(root) not in sys.path: |
| sys.path.insert(0, str(root)) |
|
|
| parser = argparse.ArgumentParser(description="Convert parquet/csv to qlib bin format") |
| parser.add_argument("--parquet", type=str, default=None, help="Source parquet path") |
| parser.add_argument("--mode", choices=["dump_all", "dump_update"], default="dump_all") |
| args = parser.parse_args() |
| run_from_parquet(args.parquet, mode=args.mode) |
|
|
