File size: 4,033 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """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)
|