SPB Git

spb/findata

Public
7.6 KB · 142 lines python
1"""Bounded historical backfill from `period=full&ticker_range=<letter>` archives → base files per (symbol, year).23Streams one vendor file (one ticker, full history) at a time so memory stays flat. Dates that already have a live4delta file are excluded (they are compacted later) so the no-duplicate invariant holds. Resumable: progress5checkpoint `job:<id>/done_symbols`; a restarted job skips symbols already published.6"""7from __future__ import annotations89from typing import Any1011import pyarrow as pa12import pyarrow.compute as pc13from findata_core.db import connection, transaction14from findata_market_firstrate import FirstRateMarketConnector, FRDNoData15from findata_schema import Interval16from findata_schema.connector import FetchRequest1718from findata_ingest import queue19from findata_ingest.lake import Lake20from findata_ingest.tasks.common import (21    dataset_for_params,22    live_delta_dates,23    record_anomalies,24    record_quality_report,25    record_raw_archive,26    upsert_coverage,27)282930def _release_memory() -> None:31    """Return freed Arrow/Polars buffers to the OS between symbols (RSS otherwise grows to several GB per worker)."""32    import contextlib33    import gc34    gc.collect()35    with contextlib.suppress(Exception):36        pa.default_memory_pool().release_unused()373839def run(job: dict[str, Any], lake: Lake, log) -> dict[str, Any]:40    p = job.get("params") or {}41    asset = p.get("type", "stock")42    letter = p.get("ticker_range") or p.get("archive_number") or p.get("update") or p.get("contract_files") or p.get("period", "full")43    timeframe = p.get("timeframe", "1min")44    adjustment = p.get("adjustment")45    only = set(s.upper() for s in (p.get("symbols") or []))46    max_symbols = int(p.get("max_symbols") or 0)47    dataset_id = job.get("dataset_id") or dataset_for_params(p)48    run_id = job.get("run_id")49    interval = Interval({"1min": "1m", "5min": "5m", "30min": "30m", "1hour": "1h", "1day": "1d"}[timeframe])50    conn_fr = FirstRateMarketConnector()51    vendor_params = {k: v for k, v in p.items() if k in ("endpoint", "type", "period", "ticker_range", "timeframe", "adjustment",52                                                           "archive_number", "update", "contract_files", "contracts", "year", "quarter") and v is not None}53    vendor_params.setdefault("endpoint", "data_file")54    if vendor_params["endpoint"] == "data_file":55        vendor_params.setdefault("period", "full")5657    ck = queue.get_checkpoint(f"job:{job['id']}", "state") or {}58    done = set(ck.get("done_symbols") or [])59    raw_rel = ck.get("raw_path")60    if raw_rel and (lake.root / raw_rel).exists():61        from datetime import datetime6263        from findata_schema.connector import FetchResult64        res = FetchResult(request=FetchRequest(dataset_id, vendor_params),65                          path=lake.abs(raw_rel), sha256=ck["sha256"], bytes=ck["bytes"], fetched_at=datetime.fromisoformat(ck["fetched_at"]),66                          vendor_last_update=None)67        log.info("backfill.resume", done=len(done), archive=raw_rel)68    else:69        req = FetchRequest(dataset_id=dataset_id, params=vendor_params)70        queue.progress(job["id"], {"stage": "download", "archive": letter})71        try:72            res = conn_fr.fetch_range(req, lake.raw_dir("firstrate", dataset_id))73        except FRDNoData as e:74            log.info("backfill.no_data", reason=str(e)[:200])75            return {"skipped": True, "reason": str(e)[:300]}76        raw_rel = lake.rel(res.path)77        with transaction() as conn:78            record_raw_archive(conn, "firstrate", dataset_id, req.params, raw_rel, res.sha256, res.bytes, res.file_count, res.vendor_last_update, run_id)79        queue.set_checkpoint(f"job:{job['id']}", "state", {"raw_path": raw_rel, "sha256": res.sha256, "bytes": res.bytes,80                                                            "fetched_at": res.fetched_at.isoformat(), "done_symbols": sorted(done)})81        log.info("backfill.downloaded", bytes=res.bytes, files=res.file_count)8283    with connection() as conn:84        delta_dates = live_delta_dates(conn, dataset_id)85    stats = {"symbols": 0, "rows": 0, "files": 0, "skipped_rows_in_delta": 0, "anomalies": 0}86    for batch in conn_fr.normalize(res, symbols=list(only) if only else None):87        if batch.symbol in done:88            continue89        if max_symbols and stats["symbols"] >= max_symbols:90            break91        rep = conn_fr.validate(batch)92        t = batch.table93        if delta_dates:94            mask = pc.invert(pc.is_in(t.column("session_date"), value_set=pa.array(sorted(delta_dates), pa.date32())))95            kept = t.filter(mask)96            stats["skipped_rows_in_delta"] += t.num_rows - kept.num_rows97            t = kept98        if t.num_rows == 0:99            done.add(batch.symbol)100            queue.set_checkpoint(f"job:{job['id']}", "state", {"raw_path": raw_rel, "sha256": res.sha256, "bytes": res.bytes,101                                                                "fetched_at": res.fetched_at.isoformat(), "done_symbols": sorted(done)})102            continue103        years = pc.year(t.column("session_date"))104        written = []105        for y in sorted(set(years.to_pylist())):106            ty = t.filter(pc.equal(years, y))107            written.append(lake.write_table(dataset_id, ty, "base", {"symbol": batch.symbol, "year": str(y)}))108        # retire previous base files of this symbol/year (re-backfill replaces)109        with connection() as conn:110            old = conn.execute("SELECT path FROM lake_files WHERE dataset_id=%s AND layout='base' AND version_removed IS NULL AND symbol_min=%s AND symbol_max=%s",111                               (dataset_id, batch.symbol, batch.symbol)).fetchall()112        lake.publish(dataset_id, written, [o["path"] for o in old], run_id, note=f"backfill {asset} {letter} {timeframe} {adjustment} {batch.symbol}")113        with transaction() as conn:114            upsert_coverage(conn, dataset_id, t, interval)115            record_anomalies(conn, dataset_id, rep.anomalies, run_id, limit=40)116        stats["symbols"] += 1117        stats["rows"] += t.num_rows118        stats["files"] += len(written)119        stats["anomalies"] += len(rep.anomalies)120        done.add(batch.symbol)121        del t, rep, written122        if stats["symbols"] % 10 == 0:123            _release_memory()124        queue.set_checkpoint(f"job:{job['id']}", "state", {"raw_path": raw_rel, "sha256": res.sha256, "bytes": res.bytes,125                                                            "fetched_at": res.fetched_at.isoformat(), "done_symbols": sorted(done)})126        queue.progress(job["id"], {"stage": "publish", **stats, "last_symbol": batch.symbol})127    with transaction() as conn:128        conn.execute("UPDATE raw_archives SET status='published' WHERE sha256=%s", (res.sha256,))129        if stats["rows"]:130            conn.execute("UPDATE datasets SET status='active', updated_at=now() WHERE id=%s AND status='planned'", (dataset_id,))131        record_quality_report(conn, dataset_id, run_id, stats["rows"], [])132    if p.get("delete_raw_after_publish", False):133        # full archives are re-downloadable; sha256 stays in raw_archives134        res.path.unlink(missing_ok=True)135        with transaction() as conn:136            conn.execute("UPDATE raw_archives SET status='deleted_after_publish', path=NULL WHERE sha256=%s", (res.sha256,))137    if interval == Interval.m1 and stats["rows"] and dataset_id in ("market_bars.stock.1m.raw", "market_bars.etf.1m.raw"):138        queue.enqueue("derive", dataset_id.replace(".1m.", ".5m."), {"source_dataset": dataset_id, "target_interval": "5m",139                      "symbols": sorted(done)}, priority=45, dedupe_key=f"derive:{dataset_id}:5m:backfill:{job['id']}")140    log.info("backfill.done", **stats)141    return stats142
Files · ae3cc68f6cf4f05be996fe98217732802cd3c838
Menu

Theme