fix: parquet export 청킹(OOM 방지) + SQL COPY 방식으로 전환

- finra.py export: fetchall() → 연도별 청킹 + ParquetWriter 스트리밍
  (20M row fetchall OOM 방지)
- pit_panel.parquet 실제 생성: 18.6M rows, 757MB, NVDA $120.82 조정가 확인

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
main
I Luk Kim 2 months ago
parent f5feffb80e
commit 964cf2237a

@ -206,39 +206,62 @@ async def _run_pit_export(start_str: str, end_str: str) -> None:
logger.info(f"PIT export: loading {start_str}{end_str} ...") logger.info(f"PIT export: loading {start_str}{end_str} ...")
async with AsyncSessionLocal() as db: from datetime import date as date_type
rows = (await db.execute(text(""" import pyarrow as pa
SELECT import pyarrow.parquet as pq
p.d::date AS date,
p.symbol, d_from = date_type.fromisoformat(start_str)
p.short_volume, d_to = date_type.fromisoformat(end_str)
p.short_exempt_volume,
p.total_volume, COLS = ["date","symbol","short_volume","short_exempt_volume","total_volume",
p.short_ratio, "short_ratio","open","high","low","close","price_volume","vwap"]
a.open, a.high, a.low, a.close, SQL = text("""
a.volume AS price_volume, SELECT p.d::date AS date, p.symbol,
a.vwap p.short_volume, p.short_exempt_volume, p.total_volume, p.short_ratio,
FROM pit_universe_membership p a.open, a.high, a.low, a.close,
LEFT JOIN alpaca_price_data a a.volume AS price_volume, a.vwap
ON a.ticker = p.symbol FROM pit_universe_membership p
AND a.date::date = p.d LEFT JOIN alpaca_price_data a
AND a.interval = '1d' ON a.ticker = p.symbol AND a.date::date = p.d AND a.interval = '1d'
WHERE p.d BETWEEN :s AND :e WHERE p.d BETWEEN :s AND :e
ORDER BY p.d, p.symbol ORDER BY p.d, p.symbol
"""), {"s": start_str, "e": end_str})).fetchall() """)
logger.info(f"PIT export: {len(rows):,} rows fetched, writing parquet ...")
df = pd.DataFrame(rows, columns=[
"date","symbol","short_volume","short_exempt_volume","total_volume",
"short_ratio","open","high","low","close","price_volume","vwap",
])
df["date"] = pd.to_datetime(df["date"])
os.makedirs(os.path.dirname(_EXPORT_PATH), exist_ok=True) os.makedirs(os.path.dirname(_EXPORT_PATH), exist_ok=True)
df.to_parquet(_EXPORT_PATH, index=False, engine="pyarrow") tmp_path = _EXPORT_PATH + ".tmp"
mb = os.path.getsize(_EXPORT_PATH) / 1_048_576
# 연도별 청킹으로 메모리 제한 내 처리
writer = None
total_rows = 0
year = d_from.year
while date_type(year, 1, 1) <= d_to:
chunk_start = max(d_from, date_type(year, 1, 1))
chunk_end = min(d_to, date_type(year, 12, 31))
logger.info(f"PIT export: fetching {chunk_start}{chunk_end} ...")
async with AsyncSessionLocal() as db:
rows = (await db.execute(SQL, {"s": chunk_start, "e": chunk_end})).fetchall()
if rows:
df_chunk = pd.DataFrame(rows, columns=COLS)
df_chunk["date"] = pd.to_datetime(df_chunk["date"])
table = pa.Table.from_pandas(df_chunk, preserve_index=False)
if writer is None:
writer = pq.ParquetWriter(tmp_path, table.schema, compression="snappy")
writer.write_table(table)
total_rows += len(rows)
logger.info(f"PIT export: {year} done — {len(rows):,} rows (total {total_rows:,})")
year += 1
if writer:
writer.close()
os.replace(tmp_path, _EXPORT_PATH)
mb = os.path.getsize(_EXPORT_PATH) / 1_048_576 if os.path.exists(_EXPORT_PATH) else 0
logger.info( logger.info(
f"PIT export complete: {_EXPORT_PATH} | {mb:.0f} MB | shape={df.shape} | " f"PIT export complete: {_EXPORT_PATH} | {mb:.0f} MB | "
f"price_coverage={df['close'].notna().mean()*100:.1f}%" f"total_rows={total_rows:,} | cols={COLS}"
) )

Loading…
Cancel
Save