35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
# sanguo_data/datawriter.py
|
|
import os
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
from vnpy.trader.object import BarData
|
|
from vnpy.trader.constant import Interval
|
|
from sanguo_data.datareader import _row_to_bar
|
|
from sanguo_data.config import DataConfig
|
|
|
|
def atomic_write_parquet(path: str, df: pd.DataFrame) -> None:
|
|
p = Path(path)
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = str(p) + ".tmp"
|
|
df.to_parquet(tmp)
|
|
os.replace(tmp, str(p)) # 原子替换
|
|
|
|
def write_daily(symbol: str, df: pd.DataFrame, cfg: DataConfig) -> None:
|
|
# 1) parquet 增量合并（按年分区，去重保留最新）
|
|
for year, group in df.groupby(df["date"].str[:4]):
|
|
f = Path(cfg.data_paths["daily_dir"]) / year / f"{symbol}.parquet"
|
|
if f.exists():
|
|
old = pd.read_parquet(f)
|
|
combined = pd.concat([old, group]).drop_duplicates("date", keep="last")
|
|
else:
|
|
combined = group
|
|
atomic_write_parquet(str(f), combined)
|
|
# 2) vnpy SQLite
|
|
bars = [_row_to_bar(symbol, row, Interval.DAILY) for _, row in df.iterrows()]
|
|
_save_to_vnpy_db(bars, cfg)
|
|
|
|
def _save_to_vnpy_db(bars: list[BarData], cfg: DataConfig) -> None:
|
|
from vnpy.trader.database import get_database
|
|
db = get_database()
|
|
db.save_bar_data(bars) # spike 验证签名
|